LLMlight 0.2.0__tar.gz → 0.3.0__tar.gz

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

@@ -18,12 +18,21 @@ import copy
 import re
 from tqdm import tqdm
 
+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 .RAG import RAG_with_RSE
+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
@@ -42,9 +51,60 @@ class LLMlight:
     5. Enable "server on local network" if you need.
     6. Enable Running.
 
+    How LLMlight Works
+    -----------------
+    LLMlight processes text through several key stages to generate intelligent responses:
+
+    1. Preprocessing Stage
+    ---------------------
+    The input context can be processed in different ways:
+    - No preprocessing: 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
+    - BERT: Advanced contextual embeddings for free-form text
+    - 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
+    - User query: The specific question or request
+    - Instructions: Additional guidance for response generation
+
+    5. Response Generation
+    ---------------------
+    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.
+    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
     ----------
-    modelname : str
+    model : str
         'hermes-3-llama-3.2-3b'
         'mistral-7b-grok'
         'openhermes-2.5-mistral-7b'
@@ -52,22 +112,29 @@ class LLMlight:
     system : str
         String of the system message.
         "I am a helpfull assistant"
-    preprocessing : str
-         None:              No pre-processing is performed. The original context is used in the pipeline of method, embedding and the response.
-        'chunk-wise':       In case you have a very large document. The text 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': In case you have a very large document. The text 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.
-    method : str
-         None:              No processing is performed. The entire context is used for the query.
-        'naive_RAG':        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':              Identify and extract entire segments of relevant text.
-    embedding : str
-        None
-        '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.
+    retrieval_method : str (default: 'RAG_basic')
+        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.
+        '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'}
+        None:                   No embedding is performed.
+        '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)
+         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.
     temperature : float, optional
-        Sampling temperature (default is 0.7).
+        Sampling temperature.
+        0.7: (default)
+        0: Deterministic
+        1: Stochastic (max)
     top_p : float, optional
         Top-p (nucleus) sampling parameter (default is 1.0, no filtering).
     chunks: dict : {'method': 'chars', 'size': 1000, 'overlap': 250, 'top_chunks': 5}
@@ -89,114 +156,77 @@ class LLMlight:
 
     Examples
     --------
-    >>> model = LLMlight()
-    >>> model.prompt('hello, who are you?')
+    >>> # Examples
+    >>> from LLMlight import LLMlight
+    >>> client =  LLMlight()
+    >>> client.prompt('hello, who are you?')
     >>> system_message = "You are a helpful assistant."
-    >>> response = model.prompt('What is the capital of France?', system=system_message, top_p=0.9)
+    >>> response = client.prompt('What is the capital of France?', system=system_message, top_p=0.9)
+    >>> print(response)
 
     """
     def __init__(self,
-                 modelname="hermes-3-llama-3.2-3b",
-                 preprocessing=None,
-                 method='naive_RAG',
-                 embedding='bert',
-                 temperature=0.7,
-                 top_p=1.0,
-                 chunks={'method': 'chars', 'size': 1000, 'overlap': 250, 'top_chunks': 5},
-                 endpoint="http://localhost:1234/v1/chat/completions",
-                 n_ctx=4096,
-                 verbose='info',
+                 model: str = "hermes-3-llama-3.2-3b",
+                 retrieval_method: str = 'RAG_basic',
+                 embedding: (str, dict) = {'memory': 'memvid', 'context': 'bert'},
+                 preprocessing: 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",
+                 n_ctx: int = 4096,
+                 verbose: (str, int) = 'info',
                  ):
 
         # Set the logger
         set_logger(verbose)
+
         # Store data in self
-        self.modelname = modelname
+        self.model = model
         self.preprocessing = preprocessing
-        self.method = method
-        self.embedding = embedding
+        self.retrieval_method = retrieval_method
+        self.alpha = alpha
+        self.top_chunks = top_chunks
         self.temperature = temperature
         self.top_p = top_p
         self.endpoint = endpoint
-        if chunks is None: chunks = {}
-        self.chunks = {**{'method': 'chars', 'size': 1000, 'overlap': 250, 'top_chunks': 5}, **chunks}
         self.n_ctx = n_ctx
         self.context = None
+        self.embedding = _set_embedding(embedding)
 
-        # Set the correct name for the model.
-        if embedding == 'bert':
-            self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
-        elif embedding == 'bge-small':
-            self.embedding_model = SentenceTransformer('BAAI/bge-small-en')
-        else:
-            self.embedding_model = None
-        # Load local model
-        if os.path.isfile(self.endpoint):
-            self.llm = load_local_gguf_model(self.endpoint, n_ctx=self.n_ctx)
-
-    def check_logger(self):
-        """Check the verbosity."""
-        logger.debug('DEBUG')
-        logger.info('INFO')
-        logger.warning('WARNING')
-        logger.critical('CRITICAL')
-
-    def get_available_models(self, validate=False):
-        # Set your local API base URL
-        # base_url = "http://localhost:1234/v1"
-        base_url = self.endpoint.split("/chat")[0]
-        logger.info('Collecting models in the API endpoint..')
-
-        # Query available models
-        response = requests.get(f"{base_url}/models")
-        model_dict = {}
-
-        try:
-            response = requests.get("http://localhost:1234/v1/models", timeout=10)
-            if response.status_code == 200:
-                try:
-                    models = response.json()["data"]
-                    model_dict = {model["id"]: model for model in models}
-                except (KeyError, ValueError) as e:
-                    logger.error("Error parsing model data:", e)
-            else:
-                logger.warning("Request failed with status code:", response.status_code)
-                logger.warning("Response:", response.text)
-
-        except requests.exceptions.RequestException as e:
-            logger.error("Request error:", e)
+        # Set chunk parameters
+        if chunks is None: chunks = {}
+        self.chunks = {**{'method': 'chars', 'size': 1000, 'overlap': 250}, **chunks}
 
-        # Check each model whether it returns a response
-        if validate:
-            logger.info("Validating models:")
-            keys = copy.deepcopy(list(model_dict.keys()))
+        # Set Memory parameters.
+        self.memory_init(retrieval_method)
 
-            for key in keys:
-                from LLMlight import LLMlight
-                llm = LLMlight(modelname=key)
-                response = llm.prompt('What is the capital of France?', system="You are a helpful assistant.", return_type='string')
-                response = response[0:30].replace('\n', ' ').replace('\r', ' ').lower()
-                if 'error: 404' in response:
-                    logger.error(f"{llm.modelname}: {response}")
-                    model_dict.pop(key)
-                else:
-                    logger.info(f"{llm.modelname}: {response}")
+        # Load local LLM gguf model
+        if os.path.isfile(self.endpoint):
+            self.llm = load_local_gguf_model(self.endpoint, n_ctx=self.n_ctx)
 
-        return list(model_dict.keys())
+        logger.info('LLMlight is initialized!')
+        logger.info(f'Model: {self.model}')
+        logger.info(f'Preprocessing: {self.preprocessing or "disabled"}')
+        logger.info(f'Retrieval method: {self.retrieval_method or "disabled"}')
+        logger.info(f'Embedding: {self.embedding or "disabled"}')
 
     def prompt(self,
-            query,
-            instructions=None,
-            system=None,
-            context=None,
-            response_format=None,
-            temperature=None,
-            top_p=None,
-            stream=False,
-            return_type='string',
-            ):
-        """
-        Run the model with the provided parameters.
+               query: str,
+               instructions: str = None,
+               system: str = None,
+               context: str = None,
+               response_format=None,
+               temperature: (int, float) = None,
+               top_p: (int, float) = None,
+               stream: bool = False,
+               return_type: str = 'string',
+               verbose=None,
+               ):
+        """Run the model with the provided parameters.
+
         The final prompt is created based on the query, instructions, and the context
 
         Parameters
@@ -225,48 +255,70 @@ class LLMlight:
             'string': Return only the string answer (remove thinking strings using tags: <think> </think>).
             'string_with_thinking' Return the full response which includes the thinking proces (if available).
 
+        Examples
+        --------
+        >>> # Examples
+        >>> from LLMlight import LLMlight
+        >>> client =  LLMlight()
+        >>> client.prompt('hello, who are you?')
+        >>> system_message = "You are a helpful assistant."
+        >>> response = client.prompt('What is the capital of France?', system=system_message, top_p=0.9)
+        >>> print(response)
+
         Returns
         -------
         str
             The model's response or an error message if the request fails.
         """
-        logger.info(f'{self.modelname} is loaded..')
-        headers = {"Content-Type": "application/json"}
+        if verbose is not None: set_logger(verbose)
+        logger.info(f'Creating response with {self.model}..')
 
+        if context is None: context = self.context
         if temperature is None: temperature = self.temperature
         if top_p is None: top_p = self.top_p
-        if context is None: context = self.context
-        # if embedding is not None: self.embedding = embedding
-        if isinstance(context, dict): context = '\n\n'.join(context.values())
-
-        # task (str): The use case for generation. (default: 'full')
-        #     'summarization'
-        #     'chat'
-        #     'code'
-        #     'longform'
-        #     'full'
-        # task = set_task(self.preprocessing, self.method)
         self.task = 'full'
 
         # Set system message
         system = set_system_message(system)
-        # Preprocessing on the context
-        processed_context = self.compute_preprocessing(query, context, instructions, system)
+
+        # Extract relevant text for video memory
+        relevant_memory = self.relevant_memory_retrieval(query, return_type='list')
+
         # Extract relevant text using retrieval method
-        relevant_context = self.relevant_text_retrieval(query, processed_context, instructions, system)
+        relevant_context = self.relevant_context_retrieval(query, context, return_type='list')
+
+        # Append the relevant chunks of texts
+        total_context = (relevant_memory or []) + (relevant_context or [])
+
+        # Preprocessing on the context
+        processed_context = self.compute_preprocessing(query, total_context, instructions, system)
+
         # Set the prompt
-        prompt = self.set_prompt(query, instructions, relevant_context, response_format=response_format)
+        logger.debug(processed_context)
+
+        # Make the prompt
+        prompt = self.set_prompt(query, instructions, processed_context, response_format=response_format)
+        logger.info('Running model now..')
 
         # Run model
         if os.path.isfile(self.endpoint):
             # Run LLM from gguf model
-            response = self.requests_post_gguf(prompt, system, temperature=temperature, top_p=top_p, headers=headers, task=self.task, stream=stream, return_type=return_type)
+            response = self.requests_post_gguf(prompt, system, temperature=temperature, top_p=top_p, task=self.task, stream=stream, return_type=return_type)
         else:
             # Run LLM with http model
-            response = self.requests_post_http(prompt, system, temperature=temperature, top_p=top_p, headers=headers, task=self.task, stream=stream, return_type=return_type)
+            response = self.requests_post_http(prompt, system, temperature=temperature, top_p=top_p, task=self.task, stream=stream, return_type=return_type)
+
         # Return
         return response
 
+    # def _get_context(self, context):
+    #     # First get the context
+    #     if context is None:
+    #         context = self.context
+    #     # if isinstance(context, dict):
+    #     #     context = '\n\n'.join(context.values())
+    #     return context
+
     def requests_post_gguf(self, prompt, system, temperature=0.8, top_p=1, headers=None, task='full', stream=False, return_type='string'):
         # Note that it is better to use messages_prompt instead of a dict (messages_dict) because most GGUF-based models don't have a tokenizer/parser that can interpret the JSON-style message structure.
         # Prepare data for request.
@@ -274,7 +326,7 @@ class LLMlight:
         # Prepare messages
         messages = [{"role": "system", "content": system}, {"role": "user", "content": prompt}]
         # Convert messages to string prompt
-        prompt = convert_prompt(messages, modelname=self.modelname)
+        prompt = convert_messages_to_model(messages, model=self.model)
         # Compute tokens
         used_tokens, max_tokens = compute_tokens(prompt, n_ctx=self.n_ctx, task=task)
 
@@ -297,19 +349,24 @@ class LLMlight:
         # Return
         return response
 
-    def requests_post_http(self, prompt, system, temperature=0.8, top_p=1, headers=None, task='full', stream=False, return_type='string'):
+    def requests_post_http(self, prompt, system, temperature=0.8, top_p=1, headers=None, task='full', stream=False, return_type='string', max_tokens=None):
         # Prepare data for request.
         if headers is None: headers = {"Content-Type": "application/json"}
         # Prepare messages
         messages = [{"role": "system", "content": system}, {"role": "user", "content": prompt}]
+
+        # Convert messages to string prompt
+        prompt = convert_messages_to_model(messages, model=self.model)
+
         # Create full prompt
         prompt = messages[0]['content'] + messages[1]['content']
         # Compute tokens
-        used_tokens, max_tokens = compute_tokens(prompt, n_ctx=self.n_ctx, task=task)
-        logger.info(f'Running {self.modelname} with max tokens: {max_tokens}')
+        if max_tokens is None:
+            used_tokens, max_tokens = compute_tokens(prompt, n_ctx=self.n_ctx, task=task)
+        # logger.info(f'Generating response with {self.model}')
 
         data = {
-            "model": self.modelname,
+            "model": self.model,
             "messages": messages,
             "temperature": temperature,
             "top_p": top_p,
@@ -350,6 +407,122 @@ 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):
+        """Build QR code video and index from chunks with unified codec handling.
+
+        Parameters
+        ----------
+        file_path : str
+            Path to output video memory file.
+        config : dict
+            Dictionary containing configuration parameters.
+
+        """
+        # Initialize video-memory
+        if not os.path.isfile(file_path): file_path = None
+        # Initialize
+        self.memory = memory.memvid_llm(file_path, config=config)
+
+        # 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)
+        # Load the retriever
+        if not hasattr(self.memory, 'retriever'):
+            # Load video-memory
+            self.memory.load()
+
+    def memory_save(self,
+                    file_path: str = None,
+                    codec: str = 'mp4v',
+                    auto_build_docker: bool = True,
+                    allow_fallback: bool = True,
+                    overwrite: bool = False,
+                    show_progress: bool = True,
+                    ):
+        """Build QR code video and index from chunks with unified codec handling.
+
+        Parameters
+        ----------
+        file_path : str (default is the initialization memory-path)
+            Path to output video memory file.
+        codec : str, optional
+            Video codec ('mp4v', 'h265', 'h264', etc.)
+            'mp4v': Default
+        auto_build_docker : bool (default: True)
+            Whether to auto-build Docker if needed.
+        allow_fallback : bool
+            Whether to fall back to MP4V if advanced codec fails.
+        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)
+
+    def memory_add(self,
+                   text: Union[str, List[str]] = None,
+                   input_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):
+        """Add chunks to memory.
+
+        Parameters
+        ----------
+        input_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)
+
+    def compute_probability(self, query, scores, embedding, n=5000):
+        if not hasattr(self.memory, 'retriever'):
+            logger.warning('No chunks to encode for null distribution. Use client.add_chunks() first')
+
+        logger.info('Creating random chunks of text with the same frequency distribution.')
+        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)))
+            bound = 'left'
+        else:
+            random_chunks = self.memory.get_random_chunks(n=n)
+            query_vector, chunk_vectors = self.fit_transform(query, random_chunks, embedding=embedding)
+            # Compute similarity
+            random_scores = cosine_similarity(query_vector, chunk_vectors)[0]
+            # Remove all scores with exactly value 0
+            random_scores = random_scores[random_scores!=0]
+            bound = 'right'
+
+        # Top indices
+        # top_indices = np.argsort(scores)[::-1]
+        # Join relevant chunks and send as prompt
+        # 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);
+
+        results = model.predict(scores, alpha=self.alpha, todf=False, multtest='fdr_bh')
+        # results['y_bool'] = results['P']<=self.alpha
+        # Store figure
+        fig, ax = model.plot(title=f'Retrieval method:{self.retrieval_method}, Embedding: {embedding}')
+
+        # Store
+        self.distfit = model
+        self.distfit.fig = fig
+        self.distfit.ax = ax
+        # Return
+        return results
+
+
+
     def task(self,
              query="Extract key insights while maintaining coherence of the previous summaries.",
              instructions="",
@@ -390,7 +563,7 @@ class LLMlight:
             logger.info(f'Working on text chunk {i}/{len(chunks)}')
 
             # Keep last N summaries for context (this needs to be within the context-window otherwise it will return an error.)
-            previous_results = "\n---\n".join(response_list[-self.chunks['top_chunks']:])
+            previous_results = "\n---\n".join(response_list[-self.top_chunks:])
 
             prompt = (
             "### Context:\n"
@@ -416,7 +589,7 @@ class LLMlight:
             response_list.append(f"Results {i+1}:\n" + chunk_result)
 
         # Final summarization pass over all collected summaries
-        results_total = "\n---\n".join(response_list[-self.chunks['top_chunks']:])
+        results_total = "\n---\n".join(response_list[-self.top_chunks:])
         final_prompt = f"""
         ### Context:
         {results_total}
@@ -443,32 +616,6 @@ class LLMlight:
         return final_result
         # return {'summary': final_result, 'summary_per_chunk': results_total}
 
-    # def query_llm(self, prompt, system=None, task='full', stream=False, return_type='string'):
-    #     """Calls the LLM and returns the response."""
-    #     # Set defaults
-    #     headers = {"Content-Type": "application/json"}
-    #     # System
-    #     if system is None: system = "You are a helpful assistant."
-    #     # Messages
-    #     messages = [{"role": "system", "content": system}, {"role": "user", "content": prompt}]
-    #     # Compute tokens
-    #     used_tokens, max_tokens = compute_tokens(prompt, n_ctx=self.n_ctx, task=task)
-
-    #     data = {
-    #         "model": self.modelname,
-    #         "messages": messages,
-    #         "temperature": self.temperature,
-    #         "top_p": self.top_p,
-    #         "stream": stream,
-    #         "max_tokens": max_tokens,
-    #         }
-
-    #     # Send POST request
-    #     response = self.requests_post(headers, data, return_type=return_type)
-
-    #     # Return
-    #     return response
-
     def global_reasoning(self, query, context, instructions, system, return_per_chunk=False, rewrite_query=False, stream=False):
         """Global Reasoning.
             1. Rewrite the input user question into something like: "Based on the extracted summaries, does the document explain the societal relevance of the research? Justify your answer."
@@ -490,19 +637,24 @@ class LLMlight:
             Only return the new question with no other information.
             """
             # Initialize model for question refinement and summarization
-            qmodel = LLMlight(modelname=self.modelname, temperature=0.7, endpoint=self.endpoint)
+            qmodel = LLMlight(model=self.model, temperature=0.7, endpoint=self.endpoint)
             # Create new query
             new_query = qmodel.prompt(query=query, instructions=instructions)
         else:
             new_query = query
 
         # Create chunks with overlapping parts to make sure we do not miss out
-        chunks = utils.chunk_text(context, method=self.chunks['method'], chunk_size=self.chunks['size'], overlap=self.chunks['overlap'])
+        if isinstance(context, str):
+            chunks = utils.chunk_text(context, method=self.chunks['method'], chunk_size=self.chunks['size'], overlap=self.chunks['overlap'])
+        else:
+            chunks = context
+
+        logger.info(f'Global-reasoning on {len(chunks)} chunks of text.')
 
         # Now summaries for the chunks
         summaries = []
-        for i, chunk in enumerate(tqdm(chunks, desc="Processing chunks", unit="chunk")):
-            logger.info(f'Working on text chunk {i+1}/{len(chunks)}')
+        for i, chunk in enumerate(tqdm(chunks, desc="Processing chunk", unit="chunk")):
+            # logger.info(f'Working on text chunk {i+1}/{len(chunks)}')
 
             prompt = f"""
             ### Context (Chunk {i+1} of {len(chunks)} from a larger document):
@@ -524,7 +676,7 @@ class LLMlight:
             """
 
             # Summarize
-            response = self.requests_post_http(prompt, system, temperature=self.temperature, top_p=self.top_p, task='summarization', stream=stream, return_type='string')
+            response = self.requests_post_http(prompt, system, temperature=self.temperature, top_p=self.top_p, task='summarization', stream=stream)
             # Append
             summaries.append(response)
             # Show
@@ -580,11 +732,16 @@ class LLMlight:
 
         """
         # Create chunks with overlapping parts to make sure we do not miss out
-        chunks = utils.chunk_text(context, method=self.chunks['method'], chunk_size=self.chunks['size'], overlap=self.chunks['overlap'])
+        if isinstance(context, str):
+            chunks = utils.chunk_text(context, method=self.chunks['method'], chunk_size=self.chunks['size'], overlap=self.chunks['overlap'], return_type='list')
+        else:
+            chunks = context
+
+        logger.info(f'Chunk wise analysis on {len(chunks)} chunks of text.')
 
         # Build a structured prompt that includes all previous summaries
         response_list = []
-        for i, chunk in enumerate(tqdm(chunks, desc="Processing chunks", unit="chunk")):
+        for i, chunk in enumerate(tqdm(chunks, desc="Processing chunk", unit="chunk")):
             logger.info(f'Working on text chunk {i+1}/{len(chunks)}')
 
             if top_chunks > 0:
@@ -703,21 +860,22 @@ class LLMlight:
         return final_response
         # return {'response': final_response, 'response_per_chunk': response_total}
 
-    def parse_large_document(self, query, context, return_type='string'):
+    def search(self, query: str, chunks: list, return_type: str = 'string', top_chunks: int = None, embedding: str = None):
         """Splits large text into chunks and finds the most relevant ones."""
-        # Create chunks
-        chunks = utils.chunk_text(context, method=self.chunks['method'], chunk_size=self.chunks['size'], overlap=self.chunks['overlap'])
         # Embedding
-        query_vector, chunk_vectors = self.fit_transform(query, chunks)
+        query_vector, chunk_vectors = self.fit_transform(query, chunks, embedding=embedding)
         # Compute similarity
-        similarities = cosine_similarity(query_vector, chunk_vectors)[0]
+        D = cosine_similarity(query_vector, chunk_vectors)[0]
         # Get top scoring chunks
-        if self.chunks['top_chunks'] is None: top_chunks = len(similarities)
-        top_indices = np.argsort(similarities)[-top_chunks:][::-1]
+        if top_chunks is None:
+            top_chunks = len(D)
+            logger.info('Number of top chunks selected is set to: {top_chunks}')
 
+        # Top indices
+        top_indices = np.argsort(D)[-top_chunks:][::-1]
         # Join relevant chunks and send as prompt
         relevant_chunks = [chunks[i] for i in top_indices]
-        relevant_scores = [similarities[i] for i in top_indices]
+        relevant_scores = [D[i] for i in top_indices]
 
         # Set the return type
         if return_type == 'score':
@@ -729,26 +887,40 @@ class LLMlight:
         else:
             return "\n---------\n".join(relevant_chunks)
 
-    def fit_transform(self, query, chunks):
+    def fit_transform(self, query, chunks, embedding=None):
         """Converts context chunks and query into vector space representations based on the selected embedding method."""
-        if self.embedding == 'tfidf':
+        if embedding is None: embedding = self.embedding
+
+        # Set embedding model parameters.
+        if self.embedding['context'] == 'bert':
+            embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+        elif self.embedding['context'] == 'bge-small':
+            embedding_model = SentenceTransformer('BAAI/bge-small-en')
+        else:
+            embedding_model = None
+
+        if embedding == 'tfidf':
             vectorizer = TfidfVectorizer()
             chunk_vectors = vectorizer.fit_transform(chunks)
             # dense_matrix = chunk_vectors.toarray()  # Converts to a NumPy array
             query_vector = vectorizer.transform([query])
-        elif self.embedding == 'bow':
+        elif embedding == 'bow':
             vectorizer = CountVectorizer()
             chunk_vectors = vectorizer.fit_transform(chunks)
             query_vector = vectorizer.transform([query])
-        # elif self.embedding_model is not None:
-        elif self.embedding == 'bert' or self.embedding == 'bge-small':
-            chunk_vectors = np.vstack([self.embedding_model.encode(chunk) for chunk in chunks])
-            query_vector = self.embedding_model.encode([query])
+        # elif embedding_model is not None:
+        elif embedding == 'bert' or embedding == 'bge-small':
+            chunk_vectors = np.vstack([embedding_model.encode(chunk) for chunk in chunks])
+            query_vector = embedding_model.encode([query])
             query_vector = query_vector.reshape(1, -1)
+        elif embedding == 'memvid':
+            logger.warning(f'Embedding method [{embedding}] can only be applied when retrieval method is set to memory-video path.')
         else:
-            raise ValueError("Unsupported embedding method. Choose a supported embedding method.")
+            logger.error(f'Available embedding methods: {get_embeddings()}')
+            raise ValueError(f"Unsupported embedding method: {self.embedding}")
+        # Return
         return query_vector, chunk_vectors
-
+        
     def compute_preprocessing(self, query, context, instructions, system):
         # Create advanced prompt using relevant chunks of text, the input query and instructions
         if context is not None:
@@ -759,7 +931,7 @@ class LLMlight:
                 # Analyze per chunk
                 relevant_context = self.chunk_wise(query, context, instructions, system, top_chunks=0, return_per_chunk=True)
             else:
-                logger.info(f'No method is applied: The entire context is used.')
+                logger.info(f'No preprocessing method is applied.')
                 relevant_context = context
         else:
             # Default
@@ -768,40 +940,107 @@ class LLMlight:
         # Return
         return relevant_context
 
-    def relevant_text_retrieval(self, query, context, instructions, system):
-        # Create advanced prompt using relevant chunks of text, the input query and instructions
+    def relevant_memory_retrieval(self, query: str, return_type='list'):
+        relevant_context = None
+
+        # Show warning if chunks are not processed yet
+        if hasattr(self, 'encoder') and len(self.encoder.chunks) > 0:
+            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):
+            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()
+
+            # Retrieval based on embedding
+            if self.embedding['memory']=='memvid':
+                # Use the memvid search retriever
+                # relevant_context = self.memory.retriever.search(query, top_k=self.top_chunks)
+                results = self.memory.retriever.index_manager.search(query, top_k=self.top_chunks)
+                scores = np.array(list(map(lambda x: x[1], results)))
+                relevant_context = list(map(lambda x: x[2]['text'], results))
+            elif self.embedding['memory'] in get_embeddings():
+                # Use the classic retrievers
+                chunks = list(map(lambda x: x.get('text'), self.memory.retriever.index_manager.metadata))
+                # Compute distances and get top k chunks
+                results = self.search(query, chunks, top_chunks=self.top_chunks, embedding=self.embedding['memory'], return_type='score')
+                scores, relevant_context = map(list, zip(*results))
+
+            # Filter on Probability
+            relevant_context = self._filter_proba(query, scores, relevant_context)
+
+            if return_type=='string':
+                # Join the chunks in context
+                relevant_context = "\n\n---\n\n".join([f"### Chunk {i+1}:\n{s}" for i, s in enumerate(relevant_context)])
+
+        # Return
+        return relevant_context
+
+    def relevant_context_retrieval(self, query, context: str, return_type='list'):
+        # Get context
+        # context = self._get_context(context)
+
         if context is not None:
-            if self.method == 'naive_RAG' and np.isin(self.embedding, ['tfidf', 'bow', 'bert', 'bge-small']):
+            # Get relevant context using RAG and embedding
+            if self.retrieval_method == 'RAG_basic' and self.embedding['context'] in get_embeddings():
                 # Find the best matching parts using simple retrieval method approach.
-                logger.info(f'[{self.method}] approach is applied with [{self.embedding}] embedding.')
-                relevant_context = self.parse_large_document(query, context, return_type='string')
-            elif self.method == 'RSE' and np.isin(self.embedding, ['bert', 'bge-small']):
-                logger.info(f'RAG approach [{self.method}] is applied.')
-                relevant_context = RAG_with_RSE(context, query, label=None, chunk_size=self.chunks['size'], irrelevant_chunk_penalty=0, embedding=self.embedding, device='cpu', batch_size=32)
+                logger.info(f"Initialize retrieval from context.. Collect [{self.top_chunks}] chunks from context using {self.embedding['context']}.")
+                # Create chunks
+                chunks = utils.chunk_text(context, method=self.chunks['method'], chunk_size=self.chunks['size'], overlap=self.chunks['overlap'])
+                # Compute distances and get top k chunks
+                results = self.search(query, chunks, top_chunks=self.top_chunks, embedding=self.embedding['context'], return_type='score')
+                # Unzip scores and context
+                scores, relevant_context = map(list, zip(*results))
+                # Filter on Probability
+                relevant_context = self._filter_proba(query, scores, relevant_context)
+
+            elif self.retrieval_method == 'RSE' and np.isin(self.embedding['context'], ['bert', 'bge-small']):
+                logger.info(f'RAG approach [{self.retrieval_method}] is applied.')
+                relevant_context = RAG.RSE(context, query, label=None, chunk_size=self.chunks['size'], irrelevant_chunk_penalty=0, embedding=self.embedding['context'], device='cpu', batch_size=32)
             else:
-                logger.info(f'No method is applied: The entire context is used.')
+                logger.info(f'No retrieval method is applied.')
                 relevant_context = context
         else:
-            # Default
             relevant_context = context
 
         # Return
         return relevant_context
 
-    def set_prompt(self, query, instructions, context, response_format=''):
+    # Compute probability
+    def _filter_proba(self, query, scores, relevant_context):
+        # Filter on significance
+        if self.alpha is not None:
+            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}')
+
+        # Return relevant context
+        return relevant_context
+
+    def set_prompt(self, query: str, instructions: str, context: (str, list), response_format: str = None):
         # Default and update when context and instructions are available.
+        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.')
+
         prompt = (
             ("Context:\n" + context + "\n\n" if context else "")
-            + (f"Instructions:\n{instructions}\n\n" if instructions != '' else "")
-            + (f"Response format:\n{response_format}\n\n" if (response_format != '') and (response_format is not None) else "")
-            + f"User question:\n"
+            + ("Instructions:\n" + instructions + "\n\n" if instructions not in ("", None) else "")
+            + ("Response format:\n" + response_format + "\n\n" if response_format not in ("", None) else "")
+            + "User question:\n"
             + query
-            )
+        )
 
         # Return
         return prompt
 
-    def read_pdf(self, filepath, 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='dict'):
         """
         Reads a PDF file and extracts its text content as a string.
 
@@ -812,26 +1051,127 @@ class LLMlight:
             str: Extracted text from the PDF.
 
         """
-        if os.path.isfile(filepath):
-            self.context = utils.read_pdf(filepath, title_pages=title_pages, body_pages=body_pages, reference_pages=reference_pages, return_type=return_type)
-            if self.context is None:
-                logger.error('No input text gathered. <return>')
-                return
-            if return_type=='dict':
-                counts = utils.count_words(self.context['body'])
-                self.context['body'] = self.context['body'] + f"\n---\nThe exact word count in this document is: {counts}"
+        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.warning(f'{filepath} does not exist.')
-            self.context = None
+            logger.error(f'file_path does not exist: {file_path}')
+
+    def get_available_models(self, validate=False):
+        """Retrieve available models from the configured API endpoint.
 
+        Optionally validates each model by sending a test prompt and filtering out
+        models that return a 404 error or similar failure response.
 
-def convert_prompt(messages, modelname='llama', add_assistant_start=True):
+        Parameters
+        ----------
+        validate : bool, optional
+            If True, each model is tested with a prompt to ensure it can respond correctly.
+            Models that fail validation (e.g., return a 404 error) are excluded from the result.
+
+        Returns
+        -------
+        list of str
+            A list of model identifiers (e.g., `"llama3"`, `"gpt-4"`) that are available and valid.
+
+        Examples
+        --------
+        >>> # Import library
+        >>> from LLMlight import LLMlight
+        <<< # Initialize
+        >>> client = LLMlight(endpoint='http://localhost:1234/v1/chat/completions', verbose='info')
+        >>> # Get models
+        >>> models = client.get_available_models(validate=False)
+        >>> # Print
+        >>> print(models)
+        >>> ['llama3', 'mistral-7b']
+
+        Notes
+        -----
+        - Requires an accessible endpoint and valid API response.
+        - 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..')
+
+        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:
+                    models = response.json()["data"]
+                    model_dict = {model["id"]: model for model in models}
+                except (KeyError, ValueError) as e:
+                    logger.error("Error parsing model data:", e)
+            else:
+                logger.warning("Request failed with status code:", response.status_code)
+                logger.warning("Response:", response.text)
+
+        except requests.exceptions.RequestException as e:
+            logger.error("Request error:", e)
+
+        # Check each model whether it returns a response
+        if validate:
+            logger.info("Validating the working of each available model. Be patient.")
+            keys = copy.deepcopy(list(model_dict.keys()))
+
+            for key in keys:
+                # logger.info(f'Checking: {key}')
+                from LLMlight import LLMlight
+                llm = LLMlight(model=key)
+                response = llm.prompt('What is the capital of France?', instructions="You are only allowed to return one word.", return_type='string')
+                response = response[0:30].replace('\n', ' ').replace('\r', ' ').lower()
+                if 'error: 404' in response:
+                    logger.error(f"{llm.model}: {response}")
+                    model_dict.pop(key)
+                else:
+                    logger.debug(f"{llm.model}: {response}")
+
+        return list(model_dict.keys())
+
+    def check_logger(self):
+        """Check the verbosity."""
+        logger.debug('DEBUG')
+        logger.info('INFO')
+        logger.warning('WARNING')
+        logger.critical('CRITICAL')
+
+#%%
+def get_embeddings():
+    return ['tfidf', 'bow', 'bert', 'bge-small', 'memvid']
+
+def _set_embedding(embedding):
+    if embedding is None: embedding = {}
+    if isinstance(embedding, str):
+        if embedding == 'automatic':
+            embedding = {'memory': 'memvid', 'context': 'bert'}
+        elif embedding in get_embeddings():
+            embedding = {'memory': embedding, 'context': embedding}
+        else:
+            embedding = {'memory': 'memvid', 'context': 'bert'}
+
+    if embedding.get('context') == 'memvid':
+        embedding['context'] = 'tfidf'
+
+    embedding = {**{'memory': 'memvid', 'context': 'tfidf'}, **embedding}
+
+    # Return
+    return embedding
+
+#%%
+def convert_messages_to_model(messages, model='llama', add_assistant_start=True):
     """
     Builds a prompt in the appropriate format for different models (LLaMA, Grok, Mistral).
 
     Args:
         messages (list of dict): Each dict must have 'role' ('system', 'user', 'assistant') and 'content'.
-        modelname (str): The type of model to generate the prompt for ('llama', 'grok', or 'mistral').
+        model (str): The type of model to generate the prompt for ('llama', 'grok', or 'mistral').
         add_assistant_start (bool): Whether to add the assistant start (default True).
         add_bos_token (bool): Helps models know it's a fresh conversation. Useful for llama/mistral/hermes-style models
 
@@ -843,31 +1183,31 @@ def convert_prompt(messages, modelname='llama', add_assistant_start=True):
         ...     {"role": "system", "content": "You are a helpful assistant."},
         ...     {"role": "user", "content": "What is the capital of France?"}
         ... ]
-        >>> prompt = convert_prompt(messages, modelname='llama')
+        >>> prompt = convert_messages_to_model(messages, model='llama')
          >>> print(prompt)
 
     """
     prompt = ""
 
-    # if add_bos_token and ('llama' in modelname or 'mistral' in modelname):
+    # if add_bos_token and ('llama' in model or 'mistral' in model):
     #     prompt += "<|begin_of_text|>\n"
 
     for msg in messages:
         role = msg["role"]
         content = msg["content"].strip()
 
-        if 'llama' in modelname or 'mistral' in modelname:
+        if 'llama' in model or 'mistral' in model:
             prompt += f"<|im_start|>{role}\n{content}\n<|im_end|>\n"
-        elif 'grok' in modelname:
+        elif 'grok' in model:
             prompt += f"<start_of_turn>{role}\n{content}<end_of_turn>\n"
         else:
             # Default to ChatML format if model not recognized
             prompt += f"<|im_start|>{role}\n{content}\n<|im_end|>\n"
 
     if add_assistant_start:
-        if 'llama' in modelname or 'mistral' in modelname:
+        if 'llama' in model or 'mistral' in model:
             prompt += "<|im_start|>assistant\n"
-        elif 'grok' in modelname:
+        elif 'grok' in model:
             prompt += "<start_of_turn>assistant\n"
 
     return prompt
@@ -923,7 +1263,9 @@ def compute_tokens(string, n_ctx=4096, task='full'):
     used_tokens = len(tokens)
     # Determine how many tokens are available for the model to generate
     max_tokens = compute_max_tokens(used_tokens, n_ctx=n_ctx, task=task)
-    logger.info(f"Used_tokens={used_tokens}, max_tokens={max_tokens}, context_limit={n_ctx}")
+    # Show message
+    logger.debug(f"Used_tokens={used_tokens}, max_tokens={max_tokens}, context_limit={n_ctx}")
+    # Return
     return used_tokens, max_tokens
 
 
@@ -962,8 +1304,19 @@ def compute_max_tokens(used_tokens, n_ctx=4096, task="full"):
 
 
 def set_system_message(system):
-    return "You are a helpful assistant." if system is None else system
+    if system is None:
+        system = """You are a helpful AI assistant with access to a knowledge base.
+
+        When answering questions:
+        1. Use the provided context from the knowledge base when relevant
+        2. When multiple sections are in the context; ### chunk 1:, ### chunk 2: or ### summary 1:, ### summary 2: etc, then the higher ranked chunks contain more relevant information.
+        3. Be clear about what information comes from the knowledge base vs. your general knowledge
+        4. If the context doesn't contain enough information, say so clearly
+        5. Provide helpful, accurate, and concise responses
 
+    The context will be provided with each query based on semantic similarity to the user's question."""
+
+    return system
 
 
 # %%
@@ -1049,11 +1402,3 @@ def disable_tqdm():
     """Set the logger for verbosity messages."""
     return (True if (logger.getEffectiveLevel()>=30) else False)
 
-
-def check_logger(verbose: [str, int] = 'info'):
-    """Check the logger."""
-    set_logger(verbose)
-    logger.debug('DEBUG')
-    logger.info('INFO')
-    logger.warning('WARNING')
-    logger.critical('CRITICAL')