LLMlight 0.1.0__py3-none-any.whl

LLMlight/RAG.py

@@ -0,0 +1,441 @@
+import pymupdf
+import os
+import numpy as np
+import json
+import re
+from sentence_transformers import SentenceTransformer
+import logging
+
+from . import utils
+# import utils
+
+logger = logging.getLogger(__name__)
+
+class SimpleVectorStore:
+    """
+    A lightweight vector store implementation using NumPy.
+    """
+    def __init__(self, dimension=1536):
+        """
+        Initialize the vector store.
+        
+        Args:
+            dimension (int): Dimension of embeddings
+        """
+        self.dimension = dimension
+        self.vectors = []
+        self.documents = []
+        self.metadata = []
+    
+    def add_documents(self, documents, vectors=None, metadata=None):
+        """
+        Add documents to the vector store.
+        
+        Args:
+            documents (List[str]): List of document chunks
+            vectors (List[List[float]], optional): List of embedding vectors
+            metadata (List[Dict], optional): List of metadata dictionaries
+        """
+        if vectors is None:
+            vectors = [None] * len(documents)
+        
+        if metadata is None:
+            metadata = [{} for _ in range(len(documents))]
+        
+        for doc, vec, meta in zip(documents, vectors, metadata):
+            self.documents.append(doc)
+            self.vectors.append(vec)
+            self.metadata.append(meta)
+    
+    def search(self, query_vector, top_k=5):
+        """
+        Search for most similar documents.
+        
+        Args:
+            query_vector (List[float]): Query embedding vector
+            top_k (int): Number of results to return
+            
+        Returns:
+            List[Dict]: List of results with documents, scores, and metadata
+        """
+        if not self.vectors or not self.documents:
+            return []
+        
+        # Convert query vector to numpy array
+        query_array = np.array(query_vector)
+        
+        # Compute similarity
+        # similarities = cosine_similarity(query_vector, chunk_vectors)[0]
+        # # Get top scoring chunks
+        # if top_chunks is None: top_chunks = len(similarities)
+        # top_indices = np.argsort(similarities)[-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]
+
+        # Calculate similarities
+        similarities = []
+        for i, vector in enumerate(self.vectors):
+            if vector is not None:
+                # Compute cosine similarity
+                similarity = np.dot(query_array, vector) / (
+                    np.linalg.norm(query_array) * np.linalg.norm(vector)
+                )
+                similarities.append((i, similarity))
+        
+        # Sort by similarity (descending)
+        similarities.sort(key=lambda x: x[1], reverse=True)
+        
+        # Get top-k results
+        results = []
+        for i, score in similarities[:top_k]:
+            results.append({
+                "document": self.documents[i],
+                "score": float(score),
+                "metadata": self.metadata[i]
+            })
+        
+        return results
+
+
+def RAG_with_RSE(context, query, label=None, chunk_size=800, irrelevant_chunk_penalty=0.2, embedding_method='bert', device='cpu', batch_size=16):
+    """
+    Complete RAG pipeline with Relevant Segment Extraction.
+
+    Args:
+        text (str):
+            Text (context)
+            Path to the PDF document
+        query (str): User query
+        chunk_size (int): Size of chunks
+        irrelevant_chunk_penalty (float): Penalty for irrelevant chunks
+        embedding_method:
+            "bge":          The large retrieval-optimized embedding model for search or ranking
+            "bge-small":    The small retrieval-optimized embedding model for search or ranking
+            "bert":         General-purpose sentence embedding model
+
+    Returns:
+        Dict: Result with query, segments, and response
+    """
+    print("\n=== STARTING RAG WITH RELEVANT SEGMENT EXTRACTION ===")
+    print(f"Query: {query}")
+
+    if embedding_method == 'bge':
+        embedding_method = "BAAI/bge-en-icl"
+    elif embedding_method == 'bge-small':
+        embedding_method = "BAAI/bge-small-en"
+    elif embedding_method == 'bert':
+        embedding_method = "all-MiniLM-L6-v2"
+    else:
+        raise Exception(f'The embedding method [{embedding_method}] is not valid.')
+
+    # Process the document to extract text, chunk it, and create embeddings
+    chunks, vector_store, doc_info = process_document(context, label=label, chunk_size=chunk_size, embedding_method=embedding_method, device=device, batch_size=batch_size)
+
+    # Calculate relevance scores and chunk values based on the query
+    print("\nCalculating relevance scores and chunk values...")
+    chunk_values = calculate_chunk_values(query, chunks, vector_store, irrelevant_chunk_penalty)
+
+    # Find the best segments of text based on chunk values
+    best_segments, scores = find_best_segments(chunk_values, max_segment_length=20, total_max_length=30, min_segment_value=0.2)
+
+    # Reconstruct text segments from the best chunks
+    print("\nReconstructing text segments from chunks...")
+    segments = reconstruct_segments(chunks, best_segments)
+
+    # Format the segments into a context string for the language model
+    context = format_segments_for_context(segments)
+
+    return context
+
+    # # Generate a response from the language model using the context
+    # response = generate_response(query, context)
+    # # Compile the result into a dictionary
+    # result = {
+    #     "query": query,
+    #     "segments": segments,
+    #     "response": response
+    # }
+    # print("\n=== FINAL RESPONSE ===")
+    # print(response)
+    # return result
+
+
+def create_embeddings(texts, model_path="BAAI/bge-small-en", batch_size=32, device="cpu"):
+    """
+    Generate embeddings using a local model.
+
+    Args:
+        texts (List[str]): List of texts to embed
+        model_path (str): Path or name of the local embedding model
+        batch_size: int : 16 when you run out of memory. Default: 32
+        embedding_method:
+            "BAAI/bge-en-icl":      The large version of retrieval-optimized embedding model for search or ranking.
+            "BAAI/bge-small-en":    The smaller version of retrieval-optimized embedding model for search or ranking.
+            "all-MiniLM-L6-v2":     general-purpose sentence embedding model
+        device: str: "cuda" or 'cpu'
+
+    Returns:
+        List[List[float]]: List of embedding vectors
+    """
+    if not texts:
+        return []
+
+    # Load the model once and reuse it
+    model = SentenceTransformer(model_path, device=device)
+
+    # Compute embeddings
+    embeddings = model.encode(texts, batch_size=batch_size, convert_to_list=True)
+    
+    # Return the list of all embeddings
+    return embeddings
+
+# def create_embeddings(texts, model="BAAI/bge-en-icl"):
+#     """
+#     Generate embeddings for texts.
+    
+#     Args:
+#         texts (List[str]): List of texts to embed
+#         model (str): Embedding model to use
+        
+#     Returns:
+#         List[List[float]]: List of embedding vectors
+#     """
+#     if not texts:
+#         return []  # Return an empty list if no texts are provided
+        
+#     # Process in batches if the list is long
+#     batch_size = 100  # Adjust based on your API limits
+#     all_embeddings = []  # Initialize a list to store all embeddings
+    
+#     for i in range(0, len(texts), batch_size):
+#         batch = texts[i:i + batch_size]  # Get the current batch of texts
+        
+#         # Create embeddings for the current batch using the specified model
+#         response = client.embeddings.create(
+#             input=batch,
+#             model=model
+#         )
+        
+#         # Extract embeddings from the response
+#         batch_embeddings = [item.embedding for item in response.data]
+#         all_embeddings.extend(batch_embeddings)  # Add the batch embeddings to the list
+        
+#     return all_embeddings  # Return the list of all embeddings
+
+def extract_text_from_pdf(pdf_path):
+    """
+    Extracts text from a PDF file and prints the first `num_chars` characters.
+
+    Args:
+    pdf_path (str): Path to the PDF file.
+
+    Returns:
+    str: Extracted text from the PDF.
+    """
+    # Open the PDF file
+    mypdf = pymupdf.open(pdf_path)
+    all_text = ""  # Initialize an empty string to store the extracted text
+
+    # Iterate through each page in the PDF
+    for page_num in range(mypdf.page_count):
+        page = mypdf[page_num]  # Get the page
+        text = page.get_text("text")  # Extract text from the page
+        all_text += text  # Append the extracted text to the all_text string
+
+    return all_text  # Return the extracted text
+
+def process_document(context, label=None, chunk_size=800, embedding_method="all-MiniLM-L6-v2", batch_size=16, device='cpu'):
+    """
+    Process a document for use with RSE.
+    
+    Args:
+        context (str): 
+            Text (context)
+            The Path to the PDF document
+        label (str):
+            Add label to identify where the specific text originates from.
+        chunk_size (int): Size of each chunk in characters
+        embedding_method:
+            "BAAI/bge-en-icl":      The large version of retrieval-optimized embedding model for search or ranking.
+            "BAAI/bge-small-en":    The smaller version of retrieval-optimized embedding model for search or ranking.
+            "all-MiniLM-L6-v2":     bert: general-purpose sentence embedding model
+
+    Returns:
+        Tuple[List[str], SimpleVectorStore, Dict]: Chunks, vector store, and document info
+    """
+    print("Extracting context from document...")
+    # Extract text from the PDF file
+    if os.path.isfile(context):
+        label = context
+        context = extract_text_from_pdf(context)
+        # text = utils.read_pdf(pdf_path, return_type='string')
+
+    print("Chunking context into non-overlapping segments...")
+    # Chunk the extracted text into non-overlapping segments
+    chunks = utils.chunk_text(context, chunk_size=chunk_size, overlap=0, method='chars')
+    print(f"Created {len(chunks)} chunks")
+
+    print("Generating embeddings for chunks...")
+    # Generate embeddings for the context chunks
+    chunk_embeddings = create_embeddings(chunks, model_path=embedding_method, batch_size=batch_size, device=device)
+
+    # Create an instance of the SimpleVectorStore
+    vector_store = SimpleVectorStore()
+
+    # Add documents with metadata (including chunk index for later reconstruction)
+    metadata = [{"chunk_index": i, "source": label} for i in range(len(chunks))]
+    vector_store.add_documents(chunks, chunk_embeddings, metadata)
+
+    # Track original document structure for segment reconstruction
+    doc_info = {
+        "chunks": chunks,
+        "source": label,
+    }
+
+    return chunks, vector_store, doc_info
+
+
+def calculate_chunk_values(query, chunks, vector_store, irrelevant_chunk_penalty=0.2):
+    """
+    Calculate chunk values by combining relevance and position.
+    
+    Args:
+        query (str): Query text
+        chunks (List[str]): List of document chunks
+        vector_store (SimpleVectorStore): Vector store containing the chunks
+        irrelevant_chunk_penalty (float): Penalty for irrelevant chunks
+        
+    Returns:
+        List[float]: List of chunk values
+    """
+    # Create query embedding
+    query_embedding = create_embeddings([query])[0]
+    
+    # Get all chunks with similarity scores
+    num_chunks = len(chunks)
+    results = vector_store.search(query_embedding, top_k=num_chunks)
+    
+    # Create a mapping of chunk_index to relevance score
+    relevance_scores = {result["metadata"]["chunk_index"]: result["score"] for result in results}
+    
+    # Calculate chunk values (relevance score minus penalty)
+    chunk_values = []
+    for i in range(num_chunks):
+        # Get relevance score or default to 0 if not in results
+        score = relevance_scores.get(i, 0.0)
+        # Apply penalty to convert to a value where irrelevant chunks have negative value
+        value = score - irrelevant_chunk_penalty
+        chunk_values.append(value)
+    
+    return chunk_values
+
+
+def find_best_segments(chunk_values, max_segment_length=20, total_max_length=30, min_segment_value=0.2):
+    """
+    Find the best segments using a variant of the maximum sum subarray algorithm.
+    
+    Args:
+        chunk_values (List[float]): Values for each chunk
+        max_segment_length (int): Maximum length of a single segment
+        total_max_length (int): Maximum total length across all segments
+        min_segment_value (float): Minimum value for a segment to be considered
+        
+    Returns:
+        List[Tuple[int, int]]: List of (start, end) indices for best segments
+    """
+    print("Finding optimal continuous text segments...")
+    
+    best_segments = []
+    segment_scores = []
+    total_included_chunks = 0
+    
+    # Keep finding segments until we hit our limits
+    while total_included_chunks < total_max_length:
+        best_score = min_segment_value  # Minimum threshold for a segment
+        best_segment = None
+        
+        # Try each possible starting position
+        for start in range(len(chunk_values)):
+            # Skip if this start position is already in a selected segment
+            if any(start >= s[0] and start < s[1] for s in best_segments):
+                continue
+                
+            # Try each possible segment length
+            for length in range(1, min(max_segment_length, len(chunk_values) - start) + 1):
+                end = start + length
+                
+                # Skip if end position is already in a selected segment
+                if any(end > s[0] and end <= s[1] for s in best_segments):
+                    continue
+                
+                # Calculate segment value as sum of chunk values
+                segment_value = sum(chunk_values[start:end])
+                
+                # Update best segment if this one is better
+                if segment_value > best_score:
+                    best_score = segment_value
+                    best_segment = (start, end)
+        
+        # If we found a good segment, add it
+        if best_segment:
+            best_segments.append(best_segment)
+            segment_scores.append(best_score)
+            total_included_chunks += best_segment[1] - best_segment[0]
+            print(f"Found segment {best_segment} with score {best_score:.4f}")
+        else:
+            # No more good segments to find
+            break
+    
+    # Sort segments by their starting position for readability
+    best_segments = sorted(best_segments, key=lambda x: x[0])
+    
+    return best_segments, segment_scores
+
+
+def reconstruct_segments(chunks, best_segments):
+    """
+    Reconstruct text segments based on chunk indices.
+    
+    Args:
+        chunks (List[str]): List of all document chunks
+        best_segments (List[Tuple[int, int]]): List of (start, end) indices for segments
+        
+    Returns:
+        List[str]: List of reconstructed text segments
+    """
+    reconstructed_segments = []  # Initialize an empty list to store the reconstructed segments
+    
+    for start, end in best_segments:
+        # Join the chunks in this segment to form the complete segment text
+        segment_text = " ".join(chunks[start:end])
+        # Append the segment text and its range to the reconstructed_segments list
+        reconstructed_segments.append({
+            "text": segment_text,
+            "segment_range": (start, end),
+        })
+    
+    return reconstructed_segments  # Return the list of reconstructed text segments
+
+
+def format_segments_for_context(segments):
+    """
+    Format segments into a context string for the LLM.
+    
+    Args:
+        segments (List[Dict]): List of segment dictionaries
+        
+    Returns:
+        str: Formatted context text
+    """
+    context = []  # Initialize an empty list to store the formatted context
+    
+    for i, segment in enumerate(segments):
+        # Create a header for each segment with its index and chunk range
+        segment_header = f"SEGMENT {i+1} (Chunks {segment['segment_range'][0]}-{segment['segment_range'][1]-1}):"
+        context.append(segment_header)  # Add the segment header to the context list
+        context.append(segment['text'])  # Add the segment text to the context list
+        context.append("-" * 80)  # Add a separator line for readability
+    
+    # Join all elements in the context list with double newlines and return the result
+    return "\n\n".join(context)

LLMlight/__init__.py

@@ -0,0 +1,37 @@
+import logging
+from LLMlight.LLMlight import LLMlight
+
+__author__ = 'Erdogan Tasksen'
+__email__ = 'erdogant@gmail.com'
+__version__ = '0.1.0'
+
+# Setup root logger
+_logger = logging.getLogger('LLMlight')
+_log_handler = logging.StreamHandler()
+_fmt = '[{asctime}] [{name}] [{levelname}] {msg}'
+_formatter = logging.Formatter(fmt=_fmt, style='{', datefmt='%d-%m-%Y %H:%M:%S')
+_log_handler.setFormatter(_formatter)
+_logger.addHandler(_log_handler)
+_log_handler.setLevel(logging.DEBUG)
+_logger.propagate = False
+
+
+
+# module level doc-string
+__doc__ = """
+LLMlight
+=====================================================================
+
+LLMlight is for...
+
+Example
+-------
+>>> import LLMlight as LLMlight
+>>> model = LLMlight.fit_transform(X)
+>>> fig,ax = LLMlight.plot(model)
+
+References
+----------
+https://github.com/erdogant/LLMlight
+
+"""

LLMlight/examples.py

@@ -0,0 +1,175 @@
+# %%
+# import LLMlight as llm
+# print(dir(llm))
+# print(llm.__version__)
+
+# %%
+from LLMlight import LLMlight
+model = LLMlight(verbose='debug')
+model.check_logger()
+
+
+#%% Available models
+from LLMlight import LLMlight
+model = LLMlight(verbose='info')
+modelnames = model.get_available_models(validate=True)
+
+# %%
+for modelname in modelnames:
+    from LLMlight import LLMlight
+    llm = LLMlight(modelname=modelname)
+    print(llm.modelname)
+
+    system_message = "You are a helpful assistant."
+    response = llm.run('What is the capital of France?', system=system_message)
+    print(response)
+
+# %%
+from LLMlight import LLMlight
+model_path = r'C:/Users\beeld/.lmstudio/models/NousResearch/Hermes-3-Llama-3.2-3B-GGUF\Hermes-3-Llama-3.2-3B.Q4_K_M.gguf'
+model = LLMlight(endpoint=model_path, top_p=0.9)
+# model.run('hello, who are you?')
+system_message = "You are a helpful assistant."
+response = model.run('What is the capital of France?', system=system_message)
+
+
+
+
+
+#%%
+import llama_cpp
+print(llama_cpp.__version__)
+print(llama_cpp.llama_cpp_version())  # Might crash if incompatible
+
+# Check your GGUF model's metadata
+model_path = r'C:/Users\beeld/.lmstudio/models/NousResearch/Hermes-3-Llama-3.2-3B-GGUF\Hermes-3-Llama-3.2-3B.Q4_K_M.gguf'
+with open(model_path, 'rb') as f:
+    header = f.read(128)
+    print(header)
+
+
+#%%
+import os
+import logging
+import requests
+from tqdm import tqdm
+from llama_cpp import Llama
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+def download_and_load_gguf_model(
+    url: str,
+    model_name: str,
+    cache_dir: str = "local_models",
+    n_ctx: int = 4096,
+    n_threads: int = 8,
+    n_gpu_layers: int = 0,
+    verbose: bool = True
+) -> Llama:
+    """
+    Downloads a GGUF model from a URL (if not already cached) and loads it with llama-cpp-python.
+
+    Args:
+        url (str): Direct URL to the .gguf model file.
+        model_name (str): Filename to use for local caching (e.g. 'Hermes.gguf').
+        cache_dir (str): Directory to store the model. Default is 'local_models'.
+        n_ctx (int): Context window size.
+        n_threads (int): CPU threads to use.
+        n_gpu_layers (int): GPU layers to offload. Use 0 for CPU-only.
+        verbose (bool): Print logs during loading.
+
+    Returns:
+        Llama: Loaded model ready for inference.
+    """
+    os.makedirs(cache_dir, exist_ok=True)
+    model_path = os.path.join(cache_dir, model_name)
+
+    if not os.path.exists(model_path):
+        logger.info(f"Model not found locally. Downloading from:\n{url}")
+        try:
+            with requests.get(url, stream=True) as r:
+                r.raise_for_status()
+                total_size = int(r.headers.get("content-length", 0))
+                with open(model_path, "wb") as f, tqdm(
+                    total=total_size, unit='B', unit_scale=True, desc=model_name
+                ) as bar:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                        bar.update(len(chunk))
+        except Exception as e:
+            raise RuntimeError(f"Failed to download model: {e}")
+
+    else:
+        logger.info(f"Using cached model at: {model_path}")
+
+    # Load with llama-cpp
+    llm = Llama(
+        model_path=model_path,
+        n_ctx=n_ctx,
+        n_threads=n_threads,
+        n_gpu_layers=n_gpu_layers,
+        verbose=verbose
+    )
+
+    logger.info("Model loaded successfully.")
+    return llm
+
+#%%
+url = "https://huggingface.co/TheBloke/Hermes-2-Pro-Llama-3-GGUF/resolve/main/hermes-2-pro-llama-3.Q4_K_M.gguf"
+
+model_name = "hermes-2-pro-llama-3.Q4_K_M.gguf"
+
+llm = download_and_load_gguf_model(url, model_name)
+
+prompt = "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\nWhat is the capital of France?\n<|start_header_id|>assistant<|end_header_id|>\n"
+response = llm(prompt=prompt, max_tokens=20, stop=["<|end_of_text|>"])
+print(response["choices"][0]["text"].strip())
+
+#%%
+# Code to inference Hermes with HF Transformers
+# Requires pytorch, transformers, bitsandbytes, sentencepiece, protobuf, and flash-attn packages
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaForCausalLM
+import bitsandbytes, flash_attn
+
+tokenizer = AutoTokenizer.from_pretrained('NousResearch/Hermes-3-Llama-3.1-8B', trust_remote_code=True)
+model = LlamaForCausalLM.from_pretrained(
+    "NousResearch/Hermes-3-Llama-3.1-8B",
+    torch_dtype=torch.float16,
+    device_map="auto",
+    load_in_8bit=False,
+    load_in_4bit=True,
+    use_flash_attention_2=True
+)
+
+prompts = [
+    """<|im_start|>system
+You are a sentient, superintelligent artificial general intelligence, here to teach and assist me.<|im_end|>
+<|im_start|>user
+Write a short story about Goku discovering kirby has teamed up with Majin Buu to destroy the world.<|im_end|>
+<|im_start|>assistant""",
+    ]
+
+for chat in prompts:
+    print(chat)
+    input_ids = tokenizer(chat, return_tensors="pt").input_ids.to("cuda")
+    generated_ids = model.generate(input_ids, max_new_tokens=750, temperature=0.8, repetition_penalty=1.1, do_sample=True, eos_token_id=tokenizer.eos_token_id)
+    response = tokenizer.decode(generated_ids[0][input_ids.shape[-1]:], skip_special_tokens=True, clean_up_tokenization_space=True)
+    print(f"Response: {response}")
+
+#%%
+url = "https://huggingface.co/TheBloke/Hermes-2-Pro-Llama-3-GGUF/resolve/main/hermes-2-pro-llama-3.Q4_K_M.gguf"
+model_name = "hermes-2-pro-llama-3.Q4_K_M.gguf"
+
+# Already avoids bitsandbytes
+llm = download_and_load_gguf_model(url, model_name, n_gpu_layers=0)
+
+
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "meta-llama/Llama-2-7b-hf"  # for example
+
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")  # will crash without CUDA if quantized