Update app.py

app.py

@@ -4,35 +4,46 @@ import pdfplumber
 import docx
 import nltk
 import gradio as gr
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import FAISS
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_text_splitters import TokenTextSplitter
-from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer
-from nltk import sent_tokenize
-from typing import List, Tuple
-from transformers import AutoModel, AutoTokenizer
-
-#import spacy
-#spacy.cli.download("en_core_web_sm")  # Ensure the model is available
-#nlp = spacy.load("en_core_web_sm")    # Load the model
-
-
+from langchain.embeddings import (
+    HuggingFaceEmbeddings,
+    OpenAIEmbeddings,
+    CohereEmbeddings,
+)
+from langchain.vectorstores import FAISS, Chroma
+from langchain.text_splitters import (
+    RecursiveCharacterTextSplitter,
+    TokenTextSplitter,
+)
+from langchain.retrievers import (
+    VectorStoreRetriever,
+    ContextualCompressionRetriever,
+)
+from langchain.retrievers.document_compressors import LLMChainExtractor
+from langchain.llms import OpenAI
+from typing import List, Dict, Any
+import pandas as pd
 
 # Ensure nltk sentence tokenizer is downloaded
-nltk.download('punkt')
+nltk.download('punkt', quiet=True)
 
 FILES_DIR = './files'
 
 # Supported embedding models
 MODELS = {
-    'e5-base': "danielheinz/e5-base-sts-en-de",
-    'multilingual-e5-base': "multilingual-e5-base",
-    'paraphrase-miniLM': "paraphrase-multilingual-MiniLM-L12-v2",
-    'paraphrase-mpnet': "paraphrase-multilingual-mpnet-base-v2",
-    'gte-large': "gte-large",
-    'gbert-base': "gbert-base"
+    'HuggingFace': {
+        'e5-base': "danielheinz/e5-base-sts-en-de",
+        'multilingual-e5-base': "multilingual-e5-base",
+        'paraphrase-miniLM': "paraphrase-multilingual-MiniLM-L12-v2",
+        'paraphrase-mpnet': "paraphrase-multilingual-mpnet-base-v2",
+        'gte-large': "gte-large",
+        'gbert-base': "gbert-base"
+    },
+    'OpenAI': {
+        'text-embedding-ada-002': "text-embedding-ada-002"
+    },
+    'Cohere': {
+        'embed-multilingual-v2.0': "embed-multilingual-v2.0"
+    }
 }
 
 class FileHandler:
@@ -63,123 +74,148 @@ class FileHandler:
         with open(file_path, 'r', encoding='utf-8') as f:
             return f.read()
 
-class EmbeddingModel:
-    def __init__(self, model_name, max_tokens=None):
-        self.model = HuggingFaceEmbeddings(model_name=model_name)
-        self.max_tokens = max_tokens
-
-    def embed(self, chunks: List[str]):
-        # Embed the list of chunks
-        return self.model.embed_documents(chunks)
-
-def process_files(model_name, split_strategy, chunk_size, overlap_size, max_tokens):
-    print('-----mmm--------') 
-    print(model_name)
-    print(split_strategy)
-    print(overlap_size)
-    print(chunk_size)
-    print(max_tokens) 
-    # File processing
-    text = ""
-    for file in os.listdir(FILES_DIR):
-        file_path = os.path.join(FILES_DIR, file)
-        text += FileHandler.extract_text(file_path)
+def get_embedding_model(model_type, model_name):
+    if model_type == 'HuggingFace':
+        return HuggingFaceEmbeddings(model_name=MODELS[model_type][model_name])
+    elif model_type == 'OpenAI':
+        return OpenAIEmbeddings(model=MODELS[model_type][model_name])
+    elif model_type == 'Cohere':
+        return CohereEmbeddings(model=MODELS[model_type][model_name])
+    else:
+        raise ValueError(f"Unsupported model type: {model_type}")
 
-    # Split text into chunks
+def get_text_splitter(split_strategy, chunk_size, overlap_size, custom_separators=None):
     if split_strategy == 'token':
-        splitter = TokenTextSplitter(chunk_size=250, chunk_overlap=20)
+        return TokenTextSplitter(chunk_size=chunk_size, chunk_overlap=overlap_size)
+    elif split_strategy == 'recursive':
+        return RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size,
+            chunk_overlap=overlap_size,
+            separators=custom_separators or ["\n\n", "\n", " ", ""]
+        )
     else:
-        splitter = RecursiveCharacterTextSplitter(chunk_size=250, chunk_overlap=20)
-    
-    chunks = splitter.split_text(text)
-    
-    # Embed chunks, not the full text
-    model = EmbeddingModel(MODELS[model_name], max_tokens=max_tokens)
-    embeddings = model.embed(chunks)
-    print(chunks)
-    return embeddings, chunks
-
-def search_embeddings(query, model_name, top_k):
-    model = HuggingFaceEmbeddings(model_name=MODELS[model_name])
-    #embeddings = model.embed_query(query)
-    embeddings = model.similarity_search(query)
-    print(embeddings[0])
-    #query = "What did the president say about Ketanji Brown Jackson"
-    #docs = db.similarity_search(query)
-    #print(docs[0].page_content)
-    # Perform FAISS or other similarity-based search over embeddings
-    # This part requires you to build and search a FAISS index with embeddings
-
-    return embeddings  # You would likely return the top-k results here
-
-def calculate_statistics(embeddings):
-    # Return time taken, token count, etc.
-    return {"tokens": len(embeddings), "time_taken": time.time()}
-
-import shutil
-
- 
-import shutil
-
-def upload_file(file, model_name, split_strategy, overlap_size,chunk_size,  max_tokens, query, top_k):
-    # Ensure chunk_size and overlap_size are valid integers and provide defaults if needed
-
-    #try:
-    #    if chunk_size is None or chunk_size == "":
-    #        chunk_size = 100  # Default value if not provided
-    #    else:
-    #        chunk_size = int(chunk_size)  # Convert to int if valid#
-
-    #    if overlap_size is None or overlap_size == "":
-    #        overlap_size = 0  # Default value if not provided
-    #    else:
-    #        overlap_size = int(overlap_size)  # Convert to int if valid
-    #except ValueError:
-    #    return {"error": "Chunk size and overlap size must be valid integers."}
-    print('-------------')
-    print(file.name)
-    print(model_name)
-    print(split_strategy)
-    print(overlap_size)
-    print(chunk_size)
-    print(max_tokens)
-    print(query)
-    print(top_k)
-      
-    # Handle file upload using Gradio file object
-    file_path = file.name  # Get the file path from Gradio file object
-
-    # Copy the uploaded file content to a local directory
-    destination_path = os.path.join(FILES_DIR, os.path.basename(file_path))
-    shutil.copyfile(file_path, destination_path)  # Use shutil to copy the file
-
-    # Process files and get embeddings
-    embeddings, chunks = process_files(model_name, split_strategy, chunk_size, overlap_size, max_tokens)
+        raise ValueError(f"Unsupported split strategy: {split_strategy}")
 
-    # Perform search
-    results = search_embeddings(query, model_name, top_k)
+def get_vector_store(store_type, texts, embedding_model):
+    if store_type == 'FAISS':
+        return FAISS.from_texts(texts, embedding_model)
+    elif store_type == 'Chroma':
+        return Chroma.from_texts(texts, embedding_model)
+    else:
+        raise ValueError(f"Unsupported vector store type: {store_type}")
 
-    # Calculate statistics
-    stats = calculate_statistics(embeddings)
+def get_retriever(vector_store, search_type, search_kwargs=None):
+    if search_type == 'similarity':
+        return vector_store.as_retriever(search_type="similarity", search_kwargs=search_kwargs)
+    elif search_type == 'mmr':
+        return vector_store.as_retriever(search_type="mmr", search_kwargs=search_kwargs)
+    else:
+        raise ValueError(f"Unsupported search type: {search_type}")
+
+def process_files(file_path, model_type, model_name, split_strategy, chunk_size, overlap_size, custom_separators):
+    # File processing
+    if file_path:
+        text = FileHandler.extract_text(file_path)
+    else:
+        text = ""
+        for file in os.listdir(FILES_DIR):
+            file_path = os.path.join(FILES_DIR, file)
+            text += FileHandler.extract_text(file_path)
+
+    # Split text into chunks
+    text_splitter = get_text_splitter(split_strategy, chunk_size, overlap_size, custom_separators)
+    chunks = text_splitter.split_text(text)
+
+    # Get embedding model
+    embedding_model = get_embedding_model(model_type, model_name)
 
-    return {"results": results, "stats": stats}
+    return chunks, embedding_model
 
+def search_embeddings(chunks, embedding_model, vector_store_type, search_type, query, top_k):
+    # Create vector store
+    vector_store = get_vector_store(vector_store_type, chunks, embedding_model)
+
+    # Get retriever
+    retriever = get_retriever(vector_store, search_type, {"k": top_k})
+
+    # Perform search
+    start_time = time.time()
+    results = retriever.get_relevant_documents(query)
+    end_time = time.time()
+
+    return results, end_time - start_time
+
+def calculate_statistics(results, search_time):
+    return {
+        "num_results": len(results),
+        "avg_content_length": sum(len(doc.page_content) for doc in results) / len(results),
+        "search_time": search_time
+    }
+
+def format_results(results, stats):
+    df = pd.DataFrame([
+        {
+            "Content": doc.page_content,
+            "Source": doc.metadata.get("source", "Unknown"),
+            "Relevance Score": doc.metadata.get("score", "N/A")
+        } for doc in results
+    ])
+    
+    formatted_stats = pd.DataFrame([stats])
+    
+    return gr.DataFrame(df), gr.DataFrame(formatted_stats)
+
+def compare_embeddings(file, query, model_types, model_names, split_strategy, chunk_size, overlap_size, custom_separators, vector_store_type, search_type, top_k):
+    all_results = []
+    all_stats = []
+
+    for model_type, model_name in zip(model_types, model_names):
+        chunks, embedding_model = process_files(
+            file.name if file else None,
+            model_type,
+            model_name,
+            split_strategy,
+            chunk_size,
+            overlap_size,
+            custom_separators.split(',') if custom_separators else None
+        )
+
+        results, search_time = search_embeddings(
+            chunks,
+            embedding_model,
+            vector_store_type,
+            search_type,
+            query,
+            top_k
+        )
+
+        stats = calculate_statistics(results, search_time)
+        stats["model"] = f"{model_type} - {model_name}"
+
+        all_results.append(results)
+        all_stats.append(stats)
+
+    return [format_results(results, stats) for results, stats in zip(all_results, all_stats)]
 
 # Gradio interface
 iface = gr.Interface(
-    fn=upload_file,
+    fn=compare_embeddings,
     inputs=[
-        gr.File(label="Upload File"),
+        gr.File(label="Upload File (Optional)"),
         gr.Textbox(label="Search Query"),
-        gr.Dropdown(choices=list(MODELS.keys()), label="Embedding Model"),
-        gr.Radio(choices=["token", "recursive"], label="Split Strategy"),
-        gr.Slider(100, 1000, step=100, value=500, label="Chunk Size"),  # Ensure type is int
-        gr.Slider(0, 100, step=10, value=50, label="Overlap Size"),  # Ensure type is int
-        gr.Slider(50, 500, step=50, value=200, label="Max Tokens"),  # Ensure type is int
-        gr.Slider(1, 10, step=1, value=5, label="Top K")  # Ensure type is int
+        gr.CheckboxGroup(choices=list(MODELS.keys()), label="Embedding Model Types", value=["HuggingFace"]),
+        gr.CheckboxGroup(choices=[model for models in MODELS.values() for model in models], label="Embedding Models", value=["e5-base"]),
+        gr.Radio(choices=["token", "recursive"], label="Split Strategy", value="recursive"),
+        gr.Slider(100, 1000, step=100, value=500, label="Chunk Size"),
+        gr.Slider(0, 100, step=10, value=50, label="Overlap Size"),
+        gr.Textbox(label="Custom Split Separators (comma-separated, optional)"),
+        gr.Radio(choices=["FAISS", "Chroma"], label="Vector Store Type", value="FAISS"),
+        gr.Radio(choices=["similarity", "mmr"], label="Search Type", value="similarity"),
+        gr.Slider(1, 10, step=1, value=5, label="Top K")
     ],
-    outputs="json"
+    outputs=[gr.DataFrame(label="Results"), gr.DataFrame(label="Statistics")] * len(MODELS),
+    title="Embedding Comparison Tool",
+    description="Compare different embedding models and retrieval strategies"
 )
 
-iface.launch()
-
+iface.launch()