|
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 |
|
|
|
|
|
|
|
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 |
|
|
|
|
|
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() |
