app.py

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