import gradio as gr
from transformers import AutoTokenizer, AutoModel
from datasets import load_dataset
import pandas as pd
import torch

# ------------------------------
# Cache model & dataset
# ------------------------------
def load_model_and_tokenizer():
    model_ckpt = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
    tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
    model = AutoModel.from_pretrained(model_ckpt)
    return tokenizer, model

def load_embeddings_dataset():
    dataset = load_dataset("ginnigarg/dataset-github-issues-embeddings", split="train")
    dataset.add_faiss_index(column="embeddings")
    return dataset

tokenizer, model = load_model_and_tokenizer()
embeddings_dataset = load_embeddings_dataset()

# ------------------------------
# Embedding helper
# ------------------------------
def cls_pooling(model_output):
    return model_output.last_hidden_state[:, 0]

def get_embeddings(text_list):
    encoded_input = tokenizer(text_list, padding=True, truncation=True, return_tensors="pt")
    with torch.no_grad():
        model_output = model(**encoded_input)
    return cls_pooling(model_output)

# ------------------------------
# Search function for Gradio
# ------------------------------
def semantic_search(question):
    question_embedding = get_embeddings([question]).cpu().detach().numpy()
    scores, samples = embeddings_dataset.get_nearest_examples(
        "embeddings", question_embedding, k=5
    )
    samples_df = pd.DataFrame.from_dict(samples)
    samples_df["scores"] = scores
    samples_df.sort_values("scores", ascending=False, inplace=True)
    # Convert to displayable format
    return samples_df[["title", "comments", "html_url", "scores"]]

# ------------------------------
# Gradio Interface
# ------------------------------
with gr.Blocks() as demo:
    gr.Markdown("# GitHub Issues Semantic Search")
    question_input = gr.Textbox(label="Enter your question", value="How can I load a dataset offline?")
    output = gr.Dataframe(headers=["title", "comments", "html_url", "scores"], datatype=["str", "str", "str", "number"])
    search_button = gr.Button("Search")

    def on_click(question):
        return semantic_search(question)

    search_button.click(fn=on_click, inputs=[question_input], outputs=[output])

# Launch the Gradio app
if __name__ == "__main__":
    demo.launch()