Update app.py

app.py

@@ -2,36 +2,35 @@ import os
 import requests
 import streamlit as st
 
-API_URL = "https://api-inference.huggingface.co/models/distilgpt2"  # Updated API endpoint
+API_URL = "https://api-inference.huggingface.co/models/gpt2"  # Updated API endpoint
 API_TOKEN = os.environ.get('HUGGINGFACEHUB_API_TOKEN')
 HEADERS = {"Authorization": f"Bearer {API_TOKEN}"}
 
-def get_sentiment_category(sentiment_label):
-    sentiment_label = sentiment_label.lower()
-    if "pos" in sentiment_label:
-        return "Positive"
-    elif "neg" in sentiment_label:
-        return "Negative"
+def get_sentiment_category(sentiment_score):
+    if sentiment_score > 0:
+        return "positive"
+    elif sentiment_score < 0:
+        return "negative"
     else:
-        return "Mixed"
+        return "mixed"
 
-st.title("DistilGPT2 Movie Review Sentiment Analysis")
+st.title("GPT-2 Movie Review Sentiment Analysis")
 
 input_text = st.text_area("Enter movie review:", "")
 
 analysis_type = st.radio("Select analysis type:", ["Zero-shot", "One-shot", "Few-shot"])
 
 if analysis_type == "Zero-shot":
-    prompt = f"Classify sentiment as positive or negative or mixed: \n\n{input_text}\n\nSentiment:"
+    prompt = f"Sentiment analysis of the following movie review: \n\n{input_text}\n\nSentiment:"
 
 elif analysis_type == "One-shot":
     example = st.text_area("Input one example:")
-    prompt = f"Classify sentiment as positive or negative or mixed: \n{example}\n\nMovie review:\n{input_text}\n\nSentiment:"
+    prompt = f"{example}\n\nMovie review:\n{input_text}\n\nSentiment:"
 
 elif analysis_type == "Few-shot":
     examples = st.text_area("Input few-shot examples, one per line:")
     examples_list = examples.split('\n')
-    prompt = f"Classify sentiment as positive or negative or mixed: \n{', '.join(examples_list)}\n\nMovie review: \n{input_text}\n\nSentiment:"
+    prompt = f"Sentiment analysis examples: \n{', '.join(examples_list)}\n\nMovie review: \n{input_text}\n\nSentiment:"
 
 if st.button("Analyze"):
     try:
@@ -40,14 +39,22 @@ if st.button("Analyze"):
 
         result = response.json()[0]['generated_text']
 
-        # Extract sentiment label directly
+        # Extract sentiment score and category
         sentiment_start = result.find("Sentiment:") + len("Sentiment:")
         sentiment_end = result.find(".", sentiment_start)
-        sentiment_label = result[sentiment_start:sentiment_end].strip()
+        sentiment_score_text = result[sentiment_start:sentiment_end].strip()
 
-        # Convert sentiment label to category
-        sentiment_category = get_sentiment_category(sentiment_label)
+        # Extract only the numerical part of the score
+        if '/' in sentiment_score_text:
+            sentiment_score_text = sentiment_score_text.split('/')[0]
+
+        try:
+            sentiment_score = float(sentiment_score_text.split()[0])
+        except ValueError:
+            sentiment_score = None
+
+        sentiment_category = get_sentiment_category(sentiment_score) if sentiment_score is not None else "mixed"
         st.write(f"Sentiment: {sentiment_category}")
 
     except requests.exceptions.RequestException as e:
-        st.error("Error reaching API\n{}".format(e))
+        st.error("Error reaching API\n{}".format(e))