Update3 app.py

app.py

@@ -1,39 +1,35 @@
 import gradio as gr
 
-import torch
-from sentence_transformers import SentenceTransformer, util
-
-import torch
+import json
 from sentence_transformers import SentenceTransformer, util
 
 def predict(body):
-    """Extracts text1 and text2 from input body, computes sentence similarity, and returns the result as a dictionary.
-
-    Args:
-        body: A dictionary containing text1 and text2 as keys.
-
-    Returns:
-        A dictionary with the key "similarity score" containing the similarity between text1 and text2 as a float.
-
-    Raises:
-        ValueError: If the input body does not contain both text1 and text2 keys.
-    """
-
-    if "text1" not in body or "text2" not in body:
-        raise ValueError("Input body must contain both 'text1' and 'text2' keys.")
-
-    t1 = body["text1"]
-    t2 = body["text2"]
-
+    # Check if body is a string and try to convert it to a dictionary
+    if isinstance(body, str):
+        try:
+            body = json.loads(body)
+        except json.JSONDecodeError:
+            return {"error": "Invalid input format. Please provide a dictionary or a valid JSON string."}
+
+    # Extract text1 and text2 from the body
+    t1 = body.get("text1")
+    t2 = body.get("text2")
+
+    # Check if text1 and text2 are not None
+    if t1 is None or t2 is None:
+        return {"error": "Missing text1 or text2 in the input."}
+
+    # Initialize the model
     model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
 
-    # Use model.encode directly for efficiency
-    embeddings = model.encode([t1, t2], convert_to_tensor=True)
-
-    similarity = util.pytorch_cos_sim(embeddings[0], embeddings[1])
-
-    return {"similarity score": float(similarity.item())}  # Extract scalar value from tensor
+    # Compute embeddings for both texts
+    embedding_1 = model.encode(t1, convert_to_tensor=True)
+    embedding_2 = model.encode(t2, convert_to_tensor=True)
 
+    # Compute the cosine similarity
+    similarity_score = util.pytorch_cos_sim(embedding_1, embedding_2)
 
+    # Return the result in the desired format
+    return {"similarity score": similarity_score.item()}
 iface = gr.Interface(fn=predict, inputs="text", outputs="text")
 iface.launch(share=True)