Update main.py

main.py

@@ -1,6 +1,4 @@
-#from fastapi import FastAPI, HTTPException, Query
-#import pandas as pd
-from fastapi import FastAPI
+from fastapi import FastAPI, Query
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 
@@ -14,32 +12,33 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 # Create a sentiment analysis pipeline
 sentiment = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)
 
+# Create a dictionary to map sentiment labels to binary values
+sentiment_label_mapping = {
+    "LABEL_1": 1,  # Positive
+    "LABEL_0": 0,  # Negative
+}
+
 # Define a request body model
 class SentimentRequest(BaseModel):
     text: str
 
 # Define a response model
 class SentimentResponse(BaseModel):
-    sentiment: str
+    sentiment: int  # 1 for positive, 0 for negative
     score: float
 
-# Create an endpoint for sentiment analysis
-@app.post("/sentiment/")
-async def analyze_sentiment(request: SentimentRequest):
-    input_text = request.text
-    result = sentiment(input_text)
+# Create an endpoint for sentiment analysis with query parameter
+@app.get("/sentiment/")
+async def analyze_sentiment(text: str = Query(..., description="Input text for sentiment analysis")):
+    result = sentiment(text)
     sentiment_label = result[0]["label"]
     sentiment_score = result[0]["score"]
-
-    if sentiment_label == "LABEL_1":
-        sentiment_label = "positive"
-    elif sentiment_label == "LABEL_0":
-        sentiment_label = "neutral"
-    else:
-        sentiment_label = "negative"
-
-    return SentimentResponse(sentiment=sentiment_label.capitalize(), score=sentiment_score)
+    
+    sentiment_value = sentiment_label_mapping.get(sentiment_label, -1)  # Default to -1 for unknown labels
+    
+    return SentimentResponse(sentiment=sentiment_value, score=sentiment_score)
 
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)
+