Update app.py

app.py

@@ -1,10 +1,7 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModel
 import torch
 import os
-
-app = FastAPI()
+import gradio as gr
 
 # Load Hugging Face Token
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -13,32 +10,35 @@ if not HF_TOKEN:
 
 # Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("mental/mental-bert-base-uncased", use_auth_token=HF_TOKEN)
-model = AutoModel.from_pretrained("mental/mental-bert-base-uncased", use_auth_token=HF_TOKEN)
+model = AutoModel.from_pretrained("mental/mental-bert-base-uncased", use_auth_token=HF_TOKEN,output_hidden_states=True)
 
 model.eval()  # Set model to evaluation mode
 
-# Request body schema
-class TextRequest(BaseModel):
-    text: str
 
-# Helper function to compute embedding
-def compute_embedding(text: str) -> list[float]:
-    """Generate a sentence embedding using mean pooling on MentalBERT output."""
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    
+def infer(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
     with torch.no_grad():
         outputs = model(**inputs)
-    embedding = outputs.last_hidden_state.mean(dim=1).squeeze()
-    return embedding.tolist()
-
-# POST endpoint to return embedding
-@app.post("/embed")
-def get_embedding(request: TextRequest):
-    text = request.text.strip()
-    if not text:
-        raise HTTPException(status_code=400, detail="Input text cannot be empty.")
     
-    try:
-        embedding = compute_embedding(text)
-        return {"embedding": embedding}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error computing embedding: {str(e)}")
+    last_hidden_state = outputs.last_hidden_state  # (1, seq_len, hidden_size)
+    mask = inputs['attention_mask'].unsqueeze(-1).expand(last_hidden_state.size()).float()
+    
+    masked_embeddings = last_hidden_state * mask
+    summed = torch.sum(masked_embeddings, dim=1)
+    counts = torch.clamp(mask.sum(dim=1), min=1e-9)
+    mean_pooled = summed / counts
+    
+    return mean_pooled.squeeze().tolist()
+
+
+# Gradio interface
+iface = gr.Interface(
+    fn=infer,
+    inputs=[
+        gr.Textbox(label="text"),
+    ],
+    outputs="text"
+)
+iface.launch()
+