Update app.py

app.py

@@ -3,17 +3,27 @@ from transformers import pipeline
 
 pipeline = pipeline("feature-extraction", model="WhereIsAI/UAE-Large-V1")
 
+
 def predict(text):
-    title_embedding = pipeline(text)
-    term_embedding = pipeline('multivitamin for men')
+    title_outputs = pipeline(text)
+    title_outputs = torch.tensor(title_outputs)
+
+    # Mean pooling
+    title_embedding = title_outputs.mean(dim=1)
+    #title_embedding = title_outputs.last_hidden_state.mean(dim=1)
+    term_outputs = pipeline('multivitamin for men')
+    term_outputs = torch.tensor(term_outputs)
+    term_embedding = term_outputs.mean(dim=1)
+    #term_embedding = term_outputs.last_hidden_state.mean(dim=1)
     semantic_score = cosine_similarity(title_embedding.flatten(), term_embedding.flatten()) * 100
     return str(format(semantic_score,'.2f'))
 
+    
 gradio_app = gr.Interface(
     predict,
     inputs='text',
     outputs='text',
-    title="Yes Or No?",
+    title="Keyword Score",
 )
 
 if __name__ == "__main__":