feat: 相似度计算demo

app.py

@@ -1,15 +1,23 @@
 import spaces
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoModel
+from numpy.linalg import norm
 
-pipe = pipeline("feature-extraction", model="jinaai/jina-embeddings-v2-base-code", trust_remote_code=True)
+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(prompt):
-    return pipe(prompt)[0].numpy().mean(axis=0)
+def generate(input1, input2):
+    embeddings = model.encode(
+        [
+            input1,
+            input2,
+        ]
+    )
+    return str(cos_sim(embeddings[0], embeddings[1]))
 
 gr.Interface(
     fn=generate,
-    inputs=gr.Text(),
-    outputs=gr.Gallery(),
+    inputs=[gr.Text(), gr.Text()],
+    outputs=gr.Text(),
 ).launch()