Spaces:
Running
Running
| import gradio as gr | |
| from sentence_transformers import SentenceTransformer | |
| import torch | |
| # ----------------------------- Load Models on Demand ----------------------------- | |
| _loaded_models = {} | |
| tab_state = "single" | |
| def load_model(model_name): | |
| if model_name not in _loaded_models: | |
| _loaded_models[model_name] = SentenceTransformer(model_name,trust_remote_code=True) | |
| return _loaded_models[model_name] | |
| # ----------------------------- Core Functions ----------------------------- | |
| def find_similar_documents(query, documents, model_name): | |
| if not query.strip(): | |
| return "Please enter a query." | |
| if not documents.strip(): | |
| return "Please enter documents (one per line)." | |
| doc_list = [d.strip() for d in documents.split('\n') if d.strip()] | |
| if not doc_list: | |
| return "Please enter at least one document." | |
| model = load_model(model_name) | |
| query_embeddings = model.encode(query, convert_to_tensor=True) | |
| doc_embeddings = model.encode(doc_list, convert_to_tensor=True) | |
| similarities = torch.nn.functional.cosine_similarity(query_embeddings, doc_embeddings) | |
| sorted_indices = torch.argsort(similarities, descending=True) | |
| results = [] | |
| for i, idx in enumerate(sorted_indices): | |
| score = similarities[idx].item() | |
| doc = doc_list[idx] | |
| results.append(f"{i+1}. Score: {score:.4f} \n  Document: {doc}") | |
| return "\n\n".join(results) | |
| def compare_models(query, documents, tarka_model, open_model): | |
| tarka_result = find_similar_documents(query, documents, tarka_model) | |
| open_result = find_similar_documents(query, documents, open_model) | |
| return tarka_result, open_result | |
| # ----------------------------- UI Layout ----------------------------- | |
| with gr.Blocks( | |
| title="Tarka Embedding Explorer", | |
| theme=gr.themes.Soft(primary_hue="blue", secondary_hue="indigo", font=["Poppins", "sans-serif"]) | |
| ) as demo: | |
| gr.Markdown(""" | |
| # 🌌 Tarka Embedding Explorer | |
| Experiment with Tarka Embedding models for semantic similarity search, or compare them with top open-source baselines. | |
| """) | |
| with gr.Tabs(): | |
| # ---------------- Single Model Tab ---------------- | |
| with gr.Tab("🔹 Single Model Search"): | |
| tab_state="single" | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| model_selector = gr.Dropdown( | |
| label="Select Model", | |
| choices=[ | |
| "Tarka-AIR/Tarka-Embedding-150M-V1", | |
| "Tarka-AIR/Tarka-Embedding-250M-V1", | |
| "Tarka-AIR/Tarka-Embedding-350M-V1", | |
| "sentence-transformers/all-MiniLM-L6-v2", | |
| "intfloat/e5-base-v2", | |
| "BAAI/bge-small-en-v1.5" | |
| ], | |
| value="Tarka-AIR/Tarka-Embedding-150M-V1" | |
| ) | |
| query_input = gr.Textbox(label="Query", placeholder="Enter your query...", lines=2) | |
| docs_input = gr.Textbox(label="Documents", placeholder="Enter one document per line...", lines=10) | |
| search_btn = gr.Button("🔍 Search", variant="primary") | |
| with gr.Column(scale=1): | |
| result_box = gr.Markdown(label="Results") | |
| search_btn.click(find_similar_documents, [query_input, docs_input, model_selector], result_box) | |
| query_input.submit(find_similar_documents, [query_input, docs_input, model_selector], result_box) | |
| examples = [ | |
| [ | |
| "Which planet is known as the Red Planet?", | |
| "Venus is often called Earth's twin because of its similar size.\nMars, known for its reddish hue, is called the Red Planet.\nJupiter, the largest planet, has a red spot.\nSaturn has iconic rings." | |
| ], | |
| [ | |
| "What causes seasons on Earth?", | |
| "The tilt of Earth's axis causes different sunlight distribution.\nThe moon affects tides but not seasons.\nEarth's orbit has minimal effect on seasons.\nRotation causes day and night." | |
| ], | |
| [ | |
| "What gas do plants release during photosynthesis?", | |
| "Plants use sunlight to convert CO₂ into glucose and release oxygen.\nAnimals inhale oxygen and exhale CO₂.\nPhotosynthesis occurs mainly in leaves." | |
| ] | |
| ] | |
| gr.Examples( | |
| examples=examples, | |
| inputs=[query_input, docs_input], | |
| label="Try Examples" | |
| ) | |
| # ---------------- Comparison Tab ---------------- | |
| with gr.Tab("⚖️ Model Comparison"): | |
| with gr.Row(): | |
| tab_state="comparsion" | |
| with gr.Column(scale=1): | |
| tarka_selector = gr.Dropdown( | |
| label="Tarka Model", | |
| choices=[ | |
| "Tarka-AIR/Tarka-Embedding-350M-V1", | |
| "Tarka-AIR/Tarka-Embedding-250M-V1", | |
| "Tarka-AIR/Tarka-Embedding-150M-V1", | |
| ], | |
| value="Tarka-AIR/Tarka-Embedding-350M-V1" | |
| ) | |
| open_selector = gr.Dropdown( | |
| label="Open-Source Model", | |
| choices=[ | |
| "sentence-transformers/all-MiniLM-L6-v2", | |
| "intfloat/e5-base-v2", | |
| "BAAI/bge-small-en-v1.5", | |
| "sentence-transformers/paraphrase-mpnet-base-v2" | |
| ], | |
| value="intfloat/e5-base-v2" | |
| ) | |
| cmp_query = gr.Textbox(label="Query", placeholder="Enter your search query...", lines=2) | |
| cmp_docs = gr.Textbox(label="Documents", placeholder="Enter one document per line...", lines=10) | |
| cmp_btn = gr.Button("⚖️ Compare Models", variant="primary") | |
| with gr.Column(scale=2): | |
| gr.Markdown("### 🧩 Comparison Results") | |
| with gr.Row(visible=False) as comparison_results: | |
| with gr.Column(): | |
| tarka_label = gr.Markdown(visible=False) | |
| tarka_output = gr.Markdown(visible=False) | |
| with gr.Column(): | |
| open_label = gr.Markdown(visible=False) | |
| open_output = gr.Markdown(visible=False) | |
| examples = [ | |
| [ | |
| "Which planet is known as the Red Planet?", | |
| "Venus is often called Earth's twin because of its similar size.\nMars, known for its reddish hue, is called the Red Planet.\nJupiter, the largest planet, has a red spot.\nSaturn has iconic rings." | |
| ], | |
| [ | |
| "What causes seasons on Earth?", | |
| "The tilt of Earth's axis causes different sunlight distribution.\nThe moon affects tides but not seasons.\nEarth's orbit has minimal effect on seasons.\nRotation causes day and night." | |
| ], | |
| [ | |
| "What gas do plants release during photosynthesis?", | |
| "Plants use sunlight to convert CO₂ into glucose and release oxygen.\nAnimals inhale oxygen and exhale CO₂.\nPhotosynthesis occurs mainly in leaves." | |
| ] | |
| ] | |
| gr.Examples( | |
| examples=examples, | |
| inputs=[cmp_query, cmp_docs], | |
| label="Try Examples" | |
| ) | |
| def run_comparison(query, docs, tarka_model, open_model): | |
| tarka_res, open_res = compare_models(query, docs, tarka_model, open_model) | |
| return ( | |
| gr.update(visible=True), | |
| gr.update(value=f"### 🧠 {tarka_model}", visible=True), | |
| gr.update(value=tarka_res, visible=True), | |
| gr.update(value=f"### 🌍 {open_model}", visible=True), | |
| gr.update(value=open_res, visible=True) | |
| ) | |
| cmp_btn.click( | |
| fn=run_comparison, | |
| inputs=[cmp_query, cmp_docs, tarka_selector, open_selector], | |
| outputs=[comparison_results, tarka_label, tarka_output, open_label, open_output] | |
| ) | |
| cmp_query.submit( | |
| fn=run_comparison, | |
| inputs=[cmp_query, cmp_docs, tarka_selector, open_selector], | |
| outputs=[comparison_results, tarka_label, tarka_output, open_label, open_output] | |
| ) | |
| # ---------------- Example Section ---------------- | |
| # Launch the app | |
| demo.launch() |