app.py

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

cos_sim = lambda a,b: (a @ b.T) / (norm(a)*norm(b))
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-code', trust_remote_code=True)

@spaces.GPU
def generate(input1, input2):
    if len(input1) < 1:
        input1 = "How do I access the index while iterating over a sequence with a for loop?"
    if len(input2) < 1:
        input2 = "# Use the built-in enumerator\nfor idx, x in enumerate(xs):\n    print(idx, x)"   
    embeddings = model.encode(
        [
            input1,
            input2,
        ]
    )
    return str(cos_sim(embeddings[0], embeddings[1]))

gr.Interface(
    fn=generate,
    inputs=[
        gr.Text(label="input1", placeholder="How do I access the index while iterating over a sequence with a for loop?"),
        gr.Text(label="input2", placeholder="# Use the built-in enumerator\nfor idx, x in enumerate(xs):\n    print(idx, x)"),
    ],
    outputs=gr.Text(),
).launch()