Update app.py

app.py

@@ -18,7 +18,129 @@ from langchain_text_splitters import (
 from typing import List, Dict, Any
 import pandas as pd
 
-# ... (previous code remains the same) ...
+
+nltk.download('punkt', quiet=True)
+
+FILES_DIR = './files'
+
+MODELS = {
+    'HuggingFace': {
+        'e5-base-de': "danielheinz/e5-base-sts-en-de",
+        '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:
+    @staticmethod
+    def extract_text(file_path):
+        ext = os.path.splitext(file_path)[-1].lower()
+        if ext == '.pdf':
+            return FileHandler._extract_from_pdf(file_path)
+        elif ext == '.docx':
+            return FileHandler._extract_from_docx(file_path)
+        elif ext == '.txt':
+            return FileHandler._extract_from_txt(file_path)
+        else:
+            raise ValueError(f"Unsupported file type: {ext}")
+
+    @staticmethod
+    def _extract_from_pdf(file_path):
+        with pdfplumber.open(file_path) as pdf:
+            return ' '.join([page.extract_text() for page in pdf.pages])
+
+    @staticmethod
+    def _extract_from_docx(file_path):
+        doc = docx.Document(file_path)
+        return ' '.join([para.text for para in doc.paragraphs])
+
+    @staticmethod
+    def _extract_from_txt(file_path):
+        with open(file_path, 'r', encoding='utf-8') as f:
+            return f.read()
+
+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}")
+
+def get_text_splitter(split_strategy, chunk_size, overlap_size, custom_separators=None):
+    if split_strategy == 'token':
+        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:
+        raise ValueError(f"Unsupported split strategy: {split_strategy}")
+
+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}")
+
+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):
+    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)
+
+    text_splitter = get_text_splitter(split_strategy, chunk_size, overlap_size, custom_separators)
+    chunks = text_splitter.split_text(text)
+
+    embedding_model = get_embedding_model(model_type, model_name)
+
+    return chunks, embedding_model, len(text.split())
+
+def search_embeddings(chunks, embedding_model, vector_store_type, search_type, query, top_k):
+    vector_store = get_vector_store(vector_store_type, chunks, embedding_model)
+    retriever = get_retriever(vector_store, search_type, {"k": top_k})
+
+    start_time = time.time()
+    results = retriever.get_relevant_documents(query)
+    end_time = time.time()
+
+    return results, end_time - start_time, vector_store
+
+def calculate_statistics(results, search_time, vector_store, num_tokens, embedding_model):
+    return {
+        "num_results": len(results),
+        "avg_content_length": sum(len(doc.page_content) for doc in results) / len(results) if results else 0,
+        "search_time": search_time,
+        "vector_store_size": vector_store._index.ntotal if hasattr(vector_store, '_index') else "N/A",
+        "num_documents": len(vector_store.docstore._dict),
+        "num_tokens": num_tokens,
+        "embedding_vocab_size": embedding_model.client.get_vocab_size() if hasattr(embedding_model, 'client') and hasattr(embedding_model.client, 'get_vocab_size') else "N/A"
+    }
 
 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 = []