AlyxTeam's picture
feat: 相似度计算demo
be10d84
raw
history blame
543 Bytes
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):
embeddings = model.encode(
[
input1,
input2,
]
)
return str(cos_sim(embeddings[0], embeddings[1]))
gr.Interface(
fn=generate,
inputs=[gr.Text(), gr.Text()],
outputs=gr.Text(),
).launch()