Update app.py

app.py

@@ -1,95 +1,31 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI
+import uvicorn
 from pydantic import BaseModel
-from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
-from dotenv import load_dotenv
-import os
+from typing import List
 import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-load_dotenv()
-hf_token = os.getenv("HF_API_TOKEN")
+class CommentsInput(BaseModel):
+    comments: List[str]
 
-model_id = "gajula21/youtube-sentiment-model-telugu"
+model_name = "gajula21/youtube-sentiment-model-telugu"  
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+model.eval() 
 
-sentiment_model = None
-label_mapping = None
+label_mapping = {0: "Negative", 1: "Neutral", 2: "Positive"}
 
-try:
-    if hf_token is None:
-        raise ValueError("HF_API_TOKEN not found in environment variables or .env file.")
-
-    device = "cpu"
-    print(f"Loading model to device: {device}")
-
-    tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token, trust_remote_code=True, use_fast=False) 
-
-    model = AutoModelForSequenceClassification.from_pretrained(
-        model_id,
-        token=hf_token,
-        device_map=device,
-        trust_remote_code=True
-    )
-
-    sentiment_model = pipeline(
-        "text-classification",
-        model=model,
-        tokenizer=tokenizer,
-    )
-
-    label_mapping = {
-        "LABEL_0": "Negative",
-        "LABEL_1": "Neutral",
-        "LABEL_2": "Positive"
-    }
-
-    print("Sentiment model loaded successfully.")
-
-except Exception as e:
-    print(f"FATAL ERROR: Failed to load sentiment model during startup: {e}")
-    raise e
-
-app = FastAPI(
-    title="YouTube Sentiment Analysis API",
-    description="API for analyzing sentiment of Telugu/English/Transliterated YouTube comments using a Hugging Face model.",
-    version="1.0.0",
-)
-
-class TextInput(BaseModel):
-    text: str
-
-class SentimentOutput(BaseModel):
-    comment: str
-    sentiment: str
-    score: float | None = None
+app = FastAPI()
 
 @app.get("/")
-async def read_root():
-    return {"message": "YouTube Sentiment Analysis API is running."}
-
-@app.post("/analyze/sentiment", response_model=SentimentOutput)
-async def analyze_single_sentiment(item: TextInput):
-    if sentiment_model is None:
-        raise HTTPException(status_code=500, detail="Sentiment model not loaded during startup.")
-
-    comment_text = item.text
-
-    if not comment_text or not comment_text.strip():
-        raise HTTPException(status_code=400, detail="Input text cannot be empty or just whitespace.")
-
-    try:
-        output = sentiment_model(comment_text)
-
-        if output and isinstance(output, list) and isinstance(output[0], dict) and 'label' in output[0]:
-            label = label_mapping.get(output[0]['label'], output[0]['label'])
-            score = output[0].get('score', None)
-
-            return SentimentOutput(
-                comment=comment_text,
-                sentiment=label,
-                score=score
-            )
-        else:
-            raise HTTPException(status_code=500, detail="Analysis output format unexpected from model pipeline.")
-
-    except Exception as e:
-        print(f"Error during sentiment analysis for text: '{comment_text[:50]}...' Error: {e}")
-        raise HTTPException(status_code=500, detail=f"Error during sentiment analysis: {e}")
+def read_root():
+    return {"message": "Hello, World!"}
+
+@app.post("/sentiment")
+def predict_sentiments(data: CommentsInput):
+    inputs = tokenizer(data.comments, return_tensors="pt", padding=True, truncation=True, max_length=256)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    predictions = torch.argmax(outputs.logits, dim=1).tolist()
+    sentiments = [label_mapping[p] for p in predictions]
+    return {"sentiments": sentiments}