Update app.py

app.py

@@ -1,33 +1,48 @@
-from transformers import AutoTokenizer, AutoModel
+import os
 import torch
 import torch.nn.functional as F
+import gradio as gr
+from transformers import AutoTokenizer, AutoModel
 
-#Mean Pooling - Take attention mask into account for correct averaging
-def mean_pooling(model_output, attention_mask):
-    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
-    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
-    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-
-
-# Sentences we want sentence embeddings for
-sentences = ['This is an example sentence', 'Each sentence is converted']
-
-# Load model from HuggingFace Hub
-tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
-model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
-
-# Tokenize sentences
-encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
-
-# Compute token embeddings
-with torch.no_grad():
-    model_output = model(**encoded_input)
-
-# Perform pooling
-sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
-
-# Normalize embeddings
-sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
+# ——— Model loading ———
+MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModel.from_pretrained(MODEL_NAME)
 
-print("Sentence embeddings:")
-print(sentence_embeddings)
+def mean_pooling(model_output, attention_mask):
+    token_embeddings = model_output[0]                        # first element: token embeddings
+    mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+    return torch.sum(token_embeddings * mask_expanded, 1) / torch.clamp(mask_expanded.sum(1), min=1e-9)
+
+# ——— Inference function ———
+def compute_similarity(sent1: str, sent2: str) -> float:
+    # Tokenize
+    encoded = tokenizer([sent1, sent2], padding=True, truncation=True, return_tensors="pt")
+    # Forward pass
+    with torch.no_grad():
+        output = model(**encoded)
+    # Pool & normalize
+    embeddings = mean_pooling(output, encoded["attention_mask"])
+    embeddings = F.normalize(embeddings, p=2, dim=1)
+    # Cosine similarity
+    sim = float((embeddings[0] @ embeddings[1]).item())
+    return round(sim, 4)
+
+# ——— Gradio interface ———
+demo = gr.Interface(
+    fn=compute_similarity,
+    inputs=[
+        gr.Textbox(label="Sentence 1", lines=2, placeholder="Enter first sentence..."),
+        gr.Textbox(label="Sentence 2", lines=2, placeholder="Enter second sentence...")
+    ],
+    outputs=gr.Number(label="Cosine Similarity"),
+    title="Sentence Embedding Similarity",
+    description="Compute the cosine similarity between the embeddings of two sentences using MiniLM."
+)
+
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=int(os.environ.get("PORT", 7860)),
+        share=False
+    )