app.py

import os
import time
import pdfplumber
import docx
import nltk
import gradio as gr
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.embeddings import (
    OpenAIEmbeddings,
    CohereEmbeddings,
)
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS, Chroma
from langchain_text_splitters import (
    RecursiveCharacterTextSplitter,
    TokenTextSplitter,
)
from typing import List, Dict, Any
import pandas as pd

# ... (previous code remains the same) ...

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 = []
    settings = {
        "split_strategy": split_strategy,
        "chunk_size": chunk_size,
        "overlap_size": overlap_size,
        "custom_separators": custom_separators,
        "vector_store_type": vector_store_type,
        "search_type": search_type,
        "top_k": top_k
    }

    for model_type, model_name in zip(model_types, model_names):
        chunks, embedding_model, num_tokens = 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, vector_store = search_embeddings(
            chunks,
            embedding_model,
            vector_store_type,
            search_type,
            query,
            top_k
        )

        stats = calculate_statistics(results, search_time, vector_store, num_tokens, embedding_model)
        stats["model"] = f"{model_type} - {model_name}"
        stats.update(settings)

        formatted_results = format_results(results, stats)
        all_results.extend(formatted_results)
        all_stats.append(stats)

    results_df = pd.DataFrame(all_results)
    stats_df = pd.DataFrame(all_stats)

    return results_df, stats_df

def format_results(results, stats):
    formatted_results = []
    for doc in results:
        result = {
            "Model": stats["model"],
            "Content": doc.page_content,
            **doc.metadata,
            **{k: v for k, v in stats.items() if k not in ["model"]}
        }
        formatted_results.append(result)
    return formatted_results

# Gradio interface
def launch_interface(share=True):
    iface = gr.Interface(
        fn=compare_embeddings,
        inputs=[
            gr.File(label="Upload File (Optional)"),
            gr.Textbox(label="Search Query"),
            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-de"]),
            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=[
            gr.Dataframe(label="Results", interactive=False),
            gr.Dataframe(label="Statistics", interactive=False)
        ],
        title="Embedding Comparison Tool",
        description="Compare different embedding models and retrieval strategies",
        examples=[
            ["example.pdf", "What is machine learning?", ["HuggingFace"], ["e5-base-de"], "recursive", 500, 50, "", "FAISS", "similarity", 5]
        ],
        allow_flagging="never"
    )

    tutorial_md = """
    # Embedding Comparison Tool Tutorial

    ... (tutorial content remains the same) ...
    """

    iface = gr.TabbedInterface(
        [iface, gr.Markdown(tutorial_md)],
        ["Embedding Comparison", "Tutorial"]
    )

    iface.launch(share=share)

if __name__ == "__main__":
    launch_interface()