import spaces
import gradio as gr
from numpy.linalg import norm
from transformers import AutoModel
from sentence_transformers import SentenceTransformer

cos_sim = lambda a,b: (a @ b.T) / (norm(a)*norm(b))

model1 = AutoModel.from_pretrained("jinaai/jina-embeddings-v2-base-code", trust_remote_code=True)
model2 = AutoModel.from_pretrained("jinaai/jina-embeddings-v2-base-en", trust_remote_code=True)
model3 = AutoModel.from_pretrained("jinaai/jina-embeddings-v2-base-zh", trust_remote_code=True)
model4 = SentenceTransformer("aspire/acge_text_embedding")
model5 = SentenceTransformer("intfloat/multilingual-e5-large")

@spaces.GPU
def generate(query1, query2, source_code):
    if len(query1) < 1:
        query1 = "How do I access the index while iterating over a sequence with a for loop?"
    if len(query2) < 1:
        query2 = "get a list of all the keys in a dictionary"
    if len(source_code) < 1:
        source_code = "# Use the built-in enumerator\nfor idx, x in enumerate(xs):\n    print(idx, x)"

    results = []
    for model in [model1, model2, model3, model4, model5]:
        embeddings = model.encode([query1, query2, source_code])
        score1 = cos_sim(embeddings[0], embeddings[2])
        score2 = cos_sim(embeddings[1], embeddings[2])
        results.append((float(score1), float(score2)))

    return results

gr.Interface(
    fn=generate,
    inputs=[
        gr.Text(label="query1", placeholder="How do I access the index while iterating over a sequence with a for loop?"),
        gr.Text(label="query2", placeholder="get a list of all the keys in a dictionary"),
        gr.Text(label="code", placeholder="# Use the built-in enumerator\nfor idx, x in enumerate(xs):\n    print(idx, x)"),
    ],
    outputs=[
        gr.Dataframe(
            headers=["Query1 Score", "Query2 Score"],
            label="Similarity Scores",
            row_labels=["jinaai/jina-embeddings-v2-base-code", "jinaai/jina-embeddings-v2-base-en", "jinaai/jina-embeddings-v2-base-zh", "aspire/acge_text_embedding", "intfloat/multilingual-e5-large"]
        )
    ],
).launch()