🔧 Updated app.py to compute embeddings using MentalBERT

.gitignore

@@ -0,0 +1,3 @@
+.env
+__pycache__/
+*.pyc

Dockerfile

@@ -1,17 +1,21 @@
-# Use an official Python image
 FROM python:3.11-slim
 
-# Set working directory
 WORKDIR /app
 
-# Copy the app files into the container
+# Copy files
 COPY . .
 
 # Install dependencies
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Expose the default port used by Hugging Face Spaces
+# Set Hugging Face cache location
+ENV HF_HOME=/app/.cache/huggingface
+ENV TRANSFORMERS_CACHE=$HF_HOME
+ENV HF_DATASETS_CACHE=$HF_HOME
+ENV HF_METRICS_CACHE=$HF_HOME
+ENV HUGGINGFACE_HUB_CACHE=$HF_HOME
+
 EXPOSE 7860
 
-# Run the FastAPI app using uvicorn
+# Run the app
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,21 +1,43 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModel
+import torch
 import os
 
 app = FastAPI()
 
-hf_token = os.getenv("HF_TOKEN")
-if not hf_token:
-    raise ValueError("HF_TOKEN not set in environment variables.")
+# Load Hugging Face Token
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("❌ Hugging Face API token not found! Set HF_TOKEN as an environment variable.")
 
-# Pass token while loading the model
-model = SentenceTransformer("mental/mental-bert-base-uncased", use_auth_token=hf_token)
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("mental/mental-bert-base-uncased", token=HF_TOKEN)
+model = AutoModel.from_pretrained("mental/mental-bert-base-uncased", token=HF_TOKEN)
+model.eval()  # Set model to evaluation mode
 
-class InputText(BaseModel):
+# 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)
+    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")
-async def embed_text(input: InputText):
-    embedding = model.encode(input.text).tolist()
-    return {"embedding": embedding}
+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)}")

requirements.txt

@@ -1,4 +1,4 @@
 fastapi
 uvicorn
-sentence-transformers
 torch
+sentence-transformers