Fix Hugging Face cache issue and add NER endpoint

Dockerfile

@@ -1,4 +1,5 @@
 FROM python:3.10-slim
+RUN mkdir -p /tmp/huggingface
 
 # Working directory
 WORKDIR /app
@@ -14,3 +15,4 @@ EXPOSE 7860
 
 # Command to run FastAPI app with uvicorn
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+

app.py

@@ -2,16 +2,21 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import pipeline
 import logging
+import os
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
-# Load models once on startup
+# Set transformers cache directory to a writable location to fix permission errors
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
+
+# Load models once on startup with error handling
 try:
     ner_model = pipeline("ner", model="dslim/bert-base-NER", aggregation_strategy="simple")
     sentiment_model = pipeline("sentiment-analysis", model="ProsusAI/finbert")
+    logger.info("Models loaded successfully.")
 except Exception as e:
     logger.error(f"Model loading failed: {e}")
     ner_model = None
@@ -28,25 +33,33 @@ def home():
 def analyze_sentiment(req: TextRequest):
     if not sentiment_model:
         raise HTTPException(status_code=503, detail="Sentiment model not available")
-    text = req.text
+    text = req.text.strip()
     if not text:
         raise HTTPException(status_code=400, detail="Text cannot be empty")
-    result = sentiment_model(text[:512])[0]
-    return {
-        "label": result["label"],
-        "score": round(result["score"] * 100, 2)
-    }
+    try:
+        result = sentiment_model(text[:512])[0]
+        return {
+            "label": result["label"],
+            "score": round(result["score"] * 100, 2)
+        }
+    except Exception as e:
+        logger.error(f"Sentiment analysis error: {e}")
+        raise HTTPException(status_code=500, detail="Sentiment analysis failed")
 
 @app.post("/ner")
 def analyze_ner(req: TextRequest):
     if not ner_model:
         raise HTTPException(status_code=503, detail="NER model not available")
-    text = req.text
+    text = req.text.strip()
     if not text:
         raise HTTPException(status_code=400, detail="Text cannot be empty")
-    entities = ner_model(text[:512])
-    # Filter relevant entities (ORG, PERSON, MISC, PRODUCT, GPE)
-    relevant = [e['word'] for e in entities if e['entity_group'] in ['ORG', 'PERSON', 'MISC', 'PRODUCT', 'GPE']]
-    # Remove duplicates and limit to 5
-    unique_entities = list(dict.fromkeys(relevant))[:5]
-    return {"entities": unique_entities}
+    try:
+        entities = ner_model(text[:512])
+        # Filter relevant entity groups
+        relevant = [e['word'] for e in entities if e.get('entity_group') in ['ORG', 'PERSON', 'MISC', 'PRODUCT', 'GPE']]
+        # Remove duplicates preserving order, limit to 5 entities
+        unique_entities = list(dict.fromkeys(relevant))[:5]
+        return {"entities": unique_entities}
+    except Exception as e:
+        logger.error(f"NER analysis error: {e}")
+        raise HTTPException(status_code=500, detail="NER analysis failed")