LLMlight 0.3.0__tar.gz → 0.4.0__tar.gz

{llmlight-0.3.0 → llmlight-0.4.0}/LLMlight/LLMlight.py

@@ -17,25 +17,27 @@ from transformers import AutoTokenizer
 import copy
 import re
 from tqdm import tqdm
+import tempfile
 
 from typing import List, Union
 
 from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 from sentence_transformers import SentenceTransformer
+from distfit import distfit
 
 from memvid import MemvidEncoder, MemvidRetriever
 # from memvid.config import get_default_config as memvid_get_default_config
 
-# import memory
-# import RAG
-# import utils
-from . import RAG
-from . import utils
-from . import memory
-# DEBUG
-# import utils
-# from RAG import RAG_with_RSE
+try:
+    from . import RAG
+    from . import utils
+    from . import memory
+except:
+    # DEBUG
+    import memory
+    import RAG
+    import utils
 
 logger = logging.getLogger(__name__)
 
@@ -52,25 +54,25 @@ class LLMlight:
     6. Enable Running.
 
     How LLMlight Works
-    -----------------
+    -------------------
     LLMlight processes text through several key stages to generate intelligent responses:
 
-    1. Preprocessing Stage
+    1. Context strategy
     ---------------------
     The input context can be processed in different ways:
-    - No preprocessing: Uses the raw context directly
+    - No Context strategy: Uses the raw context directly
     - Chunk-wise processing: Breaks down the context into manageable chunks, processes each chunk independently, and combines results
     - Global reasoning: Creates a global summary of the context before processing
 
     2. Retrieval Method Stage
-    ------------------------
+    --------------------------
     Three main approaches for retrieving relevant information:
     - Naive RAG: Splits text into chunks and uses similarity scoring to find the most relevant sections
     - RSE (Relevant Segment Extraction): Identifies and extracts complete relevant text segments
     - No retrieval: Uses the entire context directly
 
     3. Embedding Stage
-    -----------------
+    -------------------
     Multiple embedding options for text representation:
     - TF-IDF: Best for structured documents with matching query terms
     - Bag of Words: Simple word frequency approach
@@ -78,7 +80,7 @@ class LLMlight:
     - BGE-small: Efficient embedding model for general use
 
     4. Prompting Stage
-    -----------------
+    -------------------
     The system constructs prompts by combining:
     - System message: Defines the AI's role and behavior
     - Context: Processed and retrieved relevant information
@@ -86,50 +88,46 @@ class LLMlight:
     - Instructions: Additional guidance for response generation
 
     5. Response Generation
-    ---------------------
+    -----------------------
+    The system can be configured through various parameters to optimize for different use cases, from simple Q&A to complex document analysis.
     The model generates responses using:
     - Temperature control: Adjusts response randomness (0.7 default)
     - Top-p sampling: Controls response diversity
     - Context window management: Handles token limits efficiently
 
-    The system can be configured through various parameters to optimize for different use cases, from simple Q&A to complex document analysis.
-
     Processing Flow
-    --------------
-    The system follows a sequential processing flow where each stage builds upon the previous one. First, the input context undergoes preprocessing, where it can be either used as-is or transformed into chunks for more manageable processing. These chunks are then passed through the retrieval method stage, which determines how relevant information is extracted and organized.
+    ----------------
+    The system follows a sequential processing flow where each stage builds upon the previous one. First, the input context undergoes the context strategy, where it can be either used as-is or transformed into chunks for more manageable processing. These chunks are then passed through the retrieval method stage, which determines how relevant information is extracted and organized.
     During the embedding stage, the text is converted into numerical representations that capture its semantic meaning. This is crucial for the system to understand and process the content effectively. The embedding method chosen can significantly impact the system's ability to match queries with relevant content.
     The prompting stage brings together all the processed information, combining it with the user's query and any specific instructions. This creates a comprehensive prompt that guides the model in generating an appropriate response. The final response generation stage uses this prompt to create a coherent and relevant output, with parameters like temperature and top-p sampling helping to control the response's characteristics.
     Throughout this process, the system maintains flexibility through various configuration options, allowing it to adapt to different types of queries and contexts. This modular approach enables the system to handle everything from simple questions to complex document analysis tasks efficiently.
 
-
     Parameters
     ----------
     model : str
-        'hermes-3-llama-3.2-3b'
-        'mistral-7b-grok'
-        'openhermes-2.5-mistral-7b'
-        'gemma-2-9b-it'
+        'mistralai/mistral-small-3.2'
+        'qwen/qwen3-coder-30b'
+        'openai/gpt-oss-20b'
     system : str
         String of the system message.
         "I am a helpfull assistant"
-    retrieval_method : str (default: 'RAG_basic')
+    retrieval_method : str (default: 'naive_rag')
         None:                   No processing. The entire context is used for the query.
-        'knowledge_base.mp4'    Local or absolute path to your (video) memory file.
-        'RAG_basic':            Context is processed using Navie RAG approach. Ideal for chats and when you need to answer specfic questions: Chunk of text are created. Use cosine similarity to for ranking. The top scoring chunks will be combined (n chunks) and used as input with the prompt.
+        'naive_rag':            Context is processed using Navie RAG approach. Ideal for chats and when you need to answer specfic questions: Chunk of text are created. Use cosine similarity to for ranking. The top scoring chunks will be combined (n chunks) and used as input with the prompt.
         'RSE':                  Context is processed using Navie RSE approach. Identify and extract entire segments of relevant text.
     embedding : str, dict (default: 'automatic')
-        Specify the embedding. When using both video-memory and context, it can be specified with a dictionary: {'memory': 'memvid', 'RAG_basic': 'bert'}
-        'automatic':            {'memory': 'memvid', 'RAG_basic': 'bert'}
+        Specify the embedding. When using both video-memory and context, it can be specified with a dictionary: {'memory': 'memvid', 'naive_rag': 'bert'}
         None:                   No embedding is performed.
+        'automatic':            {'memory': 'memvid', 'naive_rag': 'bert'}
         'memvid':               This embedding can only be applied when using video-memory in the retrieval method.
         'tfidf':                Best use when it is a structured documents and the words in the queries are matching.
         'bow':                  Bag of words approach. Best use when you expect words in the document and queries to be matching.
         'bert':                 Best use when document is more free text and the queries may not match exactly the words or sentences in the document.
         'bge-small':
-    preprocessing : str (default: None)
+    context_strategy : str (default: None)
          None:                  No pre-processing. The original context is used in the pipeline of retrieval_method, embedding and the response.
-        'global-reasoning':     The input context will be summarized per chunk globally. The total set of summarized context is then returned. The normal pipeline proceeds for the query, instructions, system etc.
         'chunk-wise':           The input context will be analyze chunkwise based on the query, instructions and system. The total set of answered-chunks is then returned. The normal pipeline proceeds for the query, instructions, system etc.
+        'global-reasoning':     The input context will be summarized per chunk globally. The total set of summarized context is then returned. The normal pipeline proceeds for the query, instructions, system etc.
     temperature : float, optional
         Sampling temperature.
         0.7: (default)
@@ -146,13 +144,12 @@ class LLMlight:
                 With a context window (n_ctx) of 4096 your can set size=1000 words with n chunks=5 and leave some space for instructions, system and the query.
             'overlap': overlap between chunks
             'top_chunks': Retrieval of the top N chunks when performing RAG analysis.
-    endpoint : str
-        Endpoint of the LLM API
-        "http://localhost:1234/v1/chat/completions"
-        './models/Hermes-3-Llama-3.2-3B.Q4_K_M.gguf'
-        r'C:/Users/username/.lmstudio/models/lmstudio-community/gemma-2-9b-it-GGUF/gemma-2-9b-it-Q4_K_M.gguf'
     n_ctx : int, default: 4096
         The context window length is determined by the max tokens. A larger number of tokens will ask more cpu/gpu resources. Estimates: 1000 words or ~10.000 chars costs ~3000 tokens.
+    file_path : str
+        'knowledge_base.mp4'    Local or absolute path to your (video) memory file.
+    endpoint : str:  Endpoint of the LLM API
+        "http://localhost:1234/v1/chat/completions"
 
     Examples
     --------
@@ -166,26 +163,27 @@ class LLMlight:
 
     """
     def __init__(self,
-                 model: str = "hermes-3-llama-3.2-3b",
-                 retrieval_method: str = 'RAG_basic',
+                 model: str = None,
+                 retrieval_method: (None, str) = 'naive_rag',
                  embedding: (str, dict) = {'memory': 'memvid', 'context': 'bert'},
-                 preprocessing: str = None,
+                 context_strategy: str = None,
                  alpha: float = 0.05,
                  top_chunks: int = 5,
                  temperature: (int, float) = 0.7,
                  top_p: (int, float) = 1.0,
-                 chunks: dict = {'method': 'chars', 'size': 1024, 'overlap': 200},
-                 endpoint: str = "http://localhost:1234/v1/chat/completions",
+                 chunks: dict = {'method': 'chars', 'size': 1000, 'overlap': 200},
                  n_ctx: int = 4096,
+                 file_path: str = None,
+                 endpoint: str = "http://localhost:1234/v1/chat/completions",
                  verbose: (str, int) = 'info',
                  ):
 
         # Set the logger
         set_logger(verbose)
-
+        
         # Store data in self
         self.model = model
-        self.preprocessing = preprocessing
+        self.context_strategy = context_strategy
         self.retrieval_method = retrieval_method
         self.alpha = alpha
         self.top_chunks = top_chunks
@@ -195,23 +193,49 @@ class LLMlight:
         self.n_ctx = n_ctx
         self.context = None
         self.embedding = _set_embedding(embedding)
+        self.tempdir = os.path.join(tempfile.gettempdir(), 'temp_LLMlight')
+        self.file_path = self.get_full_path(file_path)
+        
+        # Make checks
+        if model is None:
+            models = self.get_available_models(validate=False)
+            logger.error(f'Model must be set.')
+            logger.info(f'Available models: {models}')
+            raise AssertionError('Model must be set.')
+
+        # Create tempdir
+        if not os.path.isdir(self.tempdir):
+            os.makedirs(self.tempdir, exist_ok=True)
 
         # Set chunk parameters
         if chunks is None: chunks = {}
         self.chunks = {**{'method': 'chars', 'size': 1000, 'overlap': 250}, **chunks}
 
         # Set Memory parameters.
-        self.memory_init(retrieval_method)
+        if file_path:
+            self.memory_load(self.file_path)
+            # self.memory_save(overwrite=False)
+
 
         # Load local LLM gguf model
         if os.path.isfile(self.endpoint):
             self.llm = load_local_gguf_model(self.endpoint, n_ctx=self.n_ctx)
 
-        logger.info('LLMlight is initialized!')
         logger.info(f'Model: {self.model}')
-        logger.info(f'Preprocessing: {self.preprocessing or "disabled"}')
+        logger.info(f'Context Strategy: {self.context_strategy or "disabled"}')
         logger.info(f'Retrieval method: {self.retrieval_method or "disabled"}')
         logger.info(f'Embedding: {self.embedding or "disabled"}')
+        logger.info('LLMlight is initialized!')
+
+    def get_full_path(self, filepath: str) -> str | None:
+        # If filepath is absolute, return as-is
+        if filepath is None:
+            return None
+        elif os.path.isabs(filepath):
+            return filepath
+        
+        # If it's relative or just a filename, prepend tempdir
+        return os.path.join(self.tempdir, filepath)
 
     def prompt(self,
                query: str,
@@ -288,17 +312,18 @@ class LLMlight:
         relevant_context = self.relevant_context_retrieval(query, context, return_type='list')
 
         # Append the relevant chunks of texts
+        if isinstance(relevant_context, str): relevant_context = [relevant_context]
         total_context = (relevant_memory or []) + (relevant_context or [])
 
-        # Preprocessing on the context
-        processed_context = self.compute_preprocessing(query, total_context, instructions, system)
+        # Context Strategu on the context
+        processed_context = self.compute_context_strategy(query, total_context, instructions, system)
 
         # Set the prompt
         logger.debug(processed_context)
 
         # Make the prompt
         prompt = self.set_prompt(query, instructions, processed_context, response_format=response_format)
-        logger.info('Running model now..')
+        logger.info(f'Running model: {self.model} ')
 
         # Run model
         if os.path.isfile(self.endpoint):
@@ -407,7 +432,7 @@ class LLMlight:
             logger.error(f"{response.status_code} - {response}")
             return f"Error: {response.status_code} - {response}"
 
-    def memory_init(self, file_path: str = "llmlight_memory.mp4", config: dict = None, embedding=None):
+    def memory_init(self, file_path: str = None, config: dict = None, embedding=None):
         """Build QR code video and index from chunks with unified codec handling.
 
         Parameters
@@ -418,31 +443,48 @@ class LLMlight:
             Dictionary containing configuration parameters.
 
         """
-        # Initialize video-memory
-        if not os.path.isfile(file_path): file_path = None
+        if file_path: self.file_path = file_path
+        # Make check whether already loaded
+        if hasattr(self, 'memory') and self.memory.file_path == self.file_path:
+            logger.info(f'Memory already initialized: [{self.file_path}] <return>')
+            return 
+
         # Initialize
         self.memory = memory.memvid_llm(file_path, config=config)
+        # Update file path
+        self.file_path = self.memory.file_path
 
         # Set the embedding
         if embedding is not None:
             self.embedding['memory'] = embedding
             logger.info(f'Memory embedding is set: {self.embedding or "disabled"}')
 
-    def memory_load(self, file_path: str = None):
-        # Initialize if not yet initialized
-        if not hasattr(self.memory, 'encoder'):
-            self.memory_init(file_path)
+    def memory_load(self, file_path: str = None, config: dict = None):
+        # Load
+        if file_path: self.file_path = file_path
+        # if not hasattr(self.memory, 'encoder'):
+        if hasattr(self, 'memory') and config is None and hasattr(self.memory, 'config'):
+            config = self.memory.config
+
+        if not hasattr(self, 'memory'):
+            self.memory = memory.memvid_llm(self.file_path, config=config)
+            self.file_path = self.memory.file_path
+            # Update file path
+            if not hasattr(self.memory, 'retriever'):
+                self.memory.load()
+            return
+
         # Load the retriever
-        if not hasattr(self.memory, 'retriever'):
-            # Load video-memory
-            self.memory.load()
+        # if not hasattr(self.memory, 'retriever'):
+        # Load video-memory retriever
+        self.memory.load()
 
     def memory_save(self,
                     file_path: str = None,
                     codec: str = 'mp4v',
-                    auto_build_docker: bool = True,
+                    auto_build_docker: bool = False,
                     allow_fallback: bool = True,
-                    overwrite: bool = False,
+                    overwrite: bool = True,
                     show_progress: bool = True,
                     ):
         """Build QR code video and index from chunks with unified codec handling.
@@ -461,31 +503,73 @@ class LLMlight:
         show_progress : bool (default: True)
 
         """
-        self.memory.save(file_path=file_path, codec=codec, auto_build_docker=auto_build_docker, allow_fallback=allow_fallback, overwrite=overwrite, show_progress=show_progress)
+        if file_path is not None: self.file_path = file_path
+        self.memory.save(file_path=self.file_path, codec=codec, auto_build_docker=auto_build_docker, allow_fallback=allow_fallback, overwrite=overwrite, show_progress=show_progress)
+        self.memory.load()
 
     def memory_add(self,
                    text: Union[str, List[str]] = None,
-                   input_files: Union[str, List[str]] = None,
+                   files: Union[str, List[str]] = None,
                    dirpath: str = None,
                    filetypes: List[str] = ['.pdf', '.txt', '.epub', '.md', '.doc', '.docx', '.rtf', '.html', '.htm'],
                    chunk_size: int = 512,
                    chunk_overlap: int = 100,
-                   overwrite=False):
+                   overwrite=True):
         """Add chunks to memory.
 
         Parameters
         ----------
-        input_files : (str, list)
+        files : (str, list)
             Path to file(s).
 
         """
-        self.memory.add(text=text, input_files=input_files, dirpath=dirpath, filetypes=filetypes, chunk_size=chunk_size, chunk_overlap=chunk_overlap, overwrite=overwrite)
+        self.memory.add(text=text, input_files=files, dirpath=dirpath, filetypes=filetypes, chunk_size=chunk_size, chunk_overlap=chunk_overlap, overwrite=overwrite, tempdir=self.tempdir)
 
-    def compute_probability(self, query, scores, embedding, n=5000):
+    def memory_chunks(self, n=10, return_type='disk'):
+        """Return the top n memory stack.
+
+        Parameters
+        ----------
+        n : int, optional
+            Top n chunks to be returned. The default is 5.
+        return_type : str, optional
+            Retrieve chunks from memory or disk. The default is 'disk'.
+            'memory'
+            'disk'
+
+        Returns
+        -------
+        chunks : list
+            Top n returned chunks.
+
+        """
+
+        if not hasattr(self, 'memory'):
+            logger.warning('Memory is not initialized. Hint: client.memory_init()')
+            return
         if not hasattr(self.memory, 'retriever'):
-            logger.warning('No chunks to encode for null distribution. Use client.add_chunks() first')
+            logger.warning('Memory is empty. Use client.memory_add() to add text and then client.memory_save() to store to disk.')
+            return
+        
+        if len(self.memory.encoder.chunks) > 0:
+            logger.warning(f'Encoder contains {len(self.memory.encoder.chunks)} chunks that are not saved yet. Save to disk with: client.memory_save()')
+            # logger.warning(f'Total chunks on disk: {len(self.memory.retriever.index_manager.metadata)}')
+            # logger.warning(f'Total chunks on memory: {len(self.memory.encoder.chunks)}')
+
+        logger.info(f'Retrieving the first top {n} chunks from {return_type}.')
+        if return_type == 'disk':
+            chunks = list(map(lambda x: x.get('text'), self.memory.retriever.index_manager.metadata))[0:n]
+        else:
+            chunks = self.memory.encoder.chunks[0:n]
 
-        logger.info('Creating random chunks of text with the same frequency distribution.')
+        return chunks
+
+    def compute_probability(self, query, scores, embedding, n=5000):
+        if not hasattr(self, 'memory') or not hasattr(self.memory, 'retriever'):
+            logger.debug('No chunks to encode for null distribution. Use client.add_chunks() first.')
+            return
+
+        logger.info('Creating null distribution -> For the detect of chunks with significant scores.')
         if self.embedding['memory']=='memvid':
             results = self.memory.retriever.index_manager.search(query, top_k=n)
             random_scores = np.array(list(map(lambda x: x[1], results)))
@@ -505,9 +589,8 @@ class LLMlight:
         # relevant_chunks = [random_chunks[i] for i in top_indices]
         # relevant_scores = [scores[i] for i in top_indices]
 
-        from distfit import distfit
         model = distfit(method='parametric', alpha=self.alpha, bound=bound, verbose='warning')
-        model.fit_transform(random_scores);
+        _ = model.fit_transform(random_scores)
 
         results = model.predict(scores, alpha=self.alpha, todf=False, multtest='fdr_bh')
         # results['y_bool'] = results['P']<=self.alpha
@@ -921,17 +1004,17 @@ class LLMlight:
         # Return
         return query_vector, chunk_vectors
         
-    def compute_preprocessing(self, query, context, instructions, system):
+    def compute_context_strategy(self, query, context, instructions, system):
         # Create advanced prompt using relevant chunks of text, the input query and instructions
         if context is not None:
-            if self.preprocessing=='global-reasoning':
+            if self.context_strategy=='global-reasoning':
                 # Global Reasoning
                 relevant_context = self.global_reasoning(query, context, instructions, system, rewrite_query=False, return_per_chunk=True)
-            elif self.preprocessing=='chunk-wise':
+            elif self.context_strategy=='chunk-wise':
                 # Analyze per chunk
                 relevant_context = self.chunk_wise(query, context, instructions, system, top_chunks=0, return_per_chunk=True)
             else:
-                logger.info(f'No preprocessing method is applied.')
+                logger.info(f'No Context Strategy method is applied.')
                 relevant_context = context
         else:
             # Default
@@ -948,11 +1031,13 @@ class LLMlight:
             logger.warning('Documents are stored in the encoder but not saved into video memory! Use save first: client.memory_save() to include the information.')
 
         # Retrieve context from video memory
-        if (self.memory.file_path is not None) and os.path.isfile(self.memory.file_path) and os.path.isfile(self.memory.index_path):
+        if self.file_path and os.path.isfile(self.file_path):
+            # and os.path.isfile(self.memory.index_path)
             logger.info(f"Initialize retrieval from memory.. Collect [{self.top_chunks}] chunks from video-memory using {self.embedding['memory']}.")
             # Initialize retriever
             # retriever = self.memory.MemvidRetriever(video_file=self.memory.file_path, index_file=self.memory['index_path'], config=self.memory.config)
-            self.memory_load()
+            if not hasattr(self, 'memory'):
+                self.memory_load()
 
             # Retrieval based on embedding
             if self.embedding['memory']=='memvid':
@@ -984,7 +1069,7 @@ class LLMlight:
 
         if context is not None:
             # Get relevant context using RAG and embedding
-            if self.retrieval_method == 'RAG_basic' and self.embedding['context'] in get_embeddings():
+            if self.retrieval_method == 'naive_rag' and self.embedding['context'] in get_embeddings():
                 # Find the best matching parts using simple retrieval method approach.
                 logger.info(f"Initialize retrieval from context.. Collect [{self.top_chunks}] chunks from context using {self.embedding['context']}.")
                 # Create chunks
@@ -1015,9 +1100,11 @@ class LLMlight:
             logger.info(f'Computing probability distribution..')
             # Compute probability
             out = self.compute_probability(query, scores, embedding=self.embedding['memory'], n=1000)
-            # Only keep significant chunks
-            relevant_context = list(np.array(relevant_context)[out['y_bool']])
-            logger.info(f'{len(relevant_context)} significant chunks found with alpha<={self.alpha}')
+            
+            if out is not None:
+                # Only keep significant chunks
+                relevant_context = list(np.array(relevant_context)[out.get('y_bool', '')])
+                logger.info(f'{len(relevant_context)} significant chunks found with alpha<={self.alpha}')
 
         # Return relevant context
         return relevant_context
@@ -1027,7 +1114,7 @@ class LLMlight:
         if isinstance(context, list):
             context = "\n\n---\n\n".join([f"### Chunk {i+1}:\n{s}" for i, s in enumerate(context)])
         if context=='':
-            logger.warning('No context is provided into the prompt.')
+            logger.info('No context is provided into the prompt.')
 
         prompt = (
             ("Context:\n" + context + "\n\n" if context else "")
@@ -1040,7 +1127,7 @@ class LLMlight:
         # Return
         return prompt
 
-    def read_pdf(self, file_path, title_pages=[1, 2], body_pages=[], reference_pages=[-1], return_type='dict'):
+    def read_pdf(self, file_path, title_pages=[1, 2], body_pages=[], reference_pages=[-1], return_type='str'):
         """
         Reads a PDF file and extracts its text content as a string.
 
@@ -1049,19 +1136,30 @@ class LLMlight:
 
         Returns:
             str: Extracted text from the PDF.
+            dict: dictionary
 
         """
+        context = ''
+
+        if 'http' in file_path[0:5]:
+            logger.info('Downloading file from url..')
+            url = file_path
+            filename = wget.filename_from_url(url)
+            file_path = os.path.join(self.tempdir, filename)
+            context = wget.download(url, file_path)
+
         if os.path.isfile(file_path):
             # Read pdf
             context = utils.read_pdf(file_path, title_pages=title_pages, body_pages=body_pages, reference_pages=reference_pages, return_type=return_type)
-            # Return
-            return context
             # if return_type=='dict':
             #     counts = utils.count_words(self.context['body'])
             #     self.context['body'] = self.context['body']
         else:
             logger.error(f'file_path does not exist: {file_path}')
 
+        # Return
+        return context
+
     def get_available_models(self, validate=False):
         """Retrieve available models from the configured API endpoint.
 
@@ -1084,7 +1182,7 @@ class LLMlight:
         >>> # Import library
         >>> from LLMlight import LLMlight
         <<< # Initialize
-        >>> client = LLMlight(endpoint='http://localhost:1234/v1/chat/completions', verbose='info')
+        >>> client = LLMlight(endpoint='http://localhost:1234/v1/chat/completions')
         >>> # Get models
         >>> models = client.get_available_models(validate=False)
         >>> # Print
@@ -1097,11 +1195,10 @@ class LLMlight:
         - Relies on the `LLMlight` class for validation (must be importable).
         """
         base_url = '/'.join(self.endpoint.split('/')[:3]) + '/'
-        logger.info('Collecting models in the API endpoint..')
+        logger.info(f'Collecting models in the API endpoint: {self.endpoint}')
 
         try:
             model_url = base_url.rstrip('/') + '/v1/models'
-            logger.info(f'Looking up models: {model_url}')
             response = requests.get(model_url, timeout=10)
             if response.status_code == 200:
                 try:
@@ -1271,16 +1368,30 @@ def compute_tokens(string, n_ctx=4096, task='full'):
 
 def compute_max_tokens(used_tokens, n_ctx=4096, task="full"):
     """
-    Compute the number of tokens that can be generated based on the task type.
+    Compute the maximum number of tokens that can be generated for a given task,
+    taking into account the number of tokens already used and the model's context window.
 
-    Parameters:
-    - used_tokens (int): Tokens already used in the input prompt.
-    - n_ctx (int): Total context window size of the model.
-    - task (str): The use case for generation. Options:
-        'summarization', 'chat', 'code', 'longform', 'full'
+    Parameters
+    ----------
+    used_tokens : int
+        Number of tokens already consumed in the current context.
+    n_ctx : int, optional
+        Total context window size of the model (default is 4096 tokens).
+    task : str, optional
+        Type of generation task. Determines the proportion of the remaining tokens to use.
+        Options are:
+        - "summarization": Use up to 50% of the context window, minimum 128 tokens.
+        - "chat": Use up to 60% of the context window, minimum 128 tokens.
+        - "code": Use up to 75% of the context window, minimum 128 tokens.
+        - "longform": Use up to 90% of the context window, minimum 256 tokens.
+        - "full": Use all remaining tokens.
+        Any unrecognized task defaults to a safe fallback using 50% of the context window.
 
-    Returns:
-    - max_tokens (int): Tokens allowed for generation.
+    Returns
+    -------
+    max_tokens : int
+        Maximum number of tokens that can be generated for the specified task,
+        ensuring at least a minimum number of tokens as defined per task type.
     """
 
     available_tokens = max(n_ctx - used_tokens, 1)  # Ensure at least 1
@@ -1402,3 +1513,32 @@ def disable_tqdm():
     """Set the logger for verbosity messages."""
     return (True if (logger.getEffectiveLevel()>=30) else False)
 
+# %% Retrieve files files.
+class wget:
+    """Retrieve file from url."""
+
+    def filename_from_url(url, ext=True):
+        """Return filename."""
+        urlname = os.path.basename(url)
+        if not ext: _, ext = os.path.splitext(urlname)
+        return urlname
+
+    def download(url, writepath):
+        """Download.
+
+        Parameters
+        ----------
+        url : str.
+            Internet source.
+        writepath : str.
+            Directory to write the file.
+
+        Returns
+        -------
+        None.
+
+        """
+        r = requests.get(url, stream=True)
+        with open(writepath, "wb") as fd:
+            for chunk in r.iter_content(chunk_size=1024):
+                fd.write(chunk)

{llmlight-0.3.0 → llmlight-0.4.0}/LLMlight/RAG.py

@@ -6,8 +6,11 @@ import re
 from sentence_transformers import SentenceTransformer
 import logging
 
-from . import utils
-# import utils
+try:
+    from . import utils
+except:
+    # DEBUG
+    import utils
 
 logger = logging.getLogger(__name__)