Spaces:
Sleeping
Sleeping
| import os | |
| import streamlit as st | |
| import google.generativeai as genai | |
| from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer | |
| import pandas as pd | |
| from db import get_mongo_client | |
| # Configure Gemini API key | |
| GEMINI_API_KEY = os.getenv("gemini_api") | |
| if GEMINI_API_KEY: | |
| genai.configure(api_key=GEMINI_API_KEY) | |
| else: | |
| st.error("⚠️ Google API key is missing! Set it in Hugging Face Secrets.") | |
| # Load pre-trained sentiment analysis model | |
| MODEL_NAME = "cardiffnlp/twitter-roberta-base-sentiment" | |
| try: | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
| model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME) | |
| sentiment_pipeline = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer) | |
| except Exception as e: | |
| st.error(f"❌ Error loading sentiment model: {e}") | |
| # Load Topic Extraction Model | |
| try: | |
| topic_pipeline = pipeline("zero-shot-classification", model="facebook/bart-large-mnli") | |
| except Exception as e: | |
| st.error(f"❌ Error loading topic extraction model: {e}") | |
| # Predefined topic labels for classification | |
| TOPIC_LABELS = [ | |
| "Technology", "Politics", "Business", "Sports", "Entertainment", | |
| "Health", "Science", "Education", "Finance", "Travel", "Food" | |
| ] | |
| # Function to analyze sentiment using the pre-trained model | |
| def analyze_sentiment(text): | |
| try: | |
| sentiment_result = sentiment_pipeline(text)[0] | |
| label = sentiment_result['label'] | |
| score = sentiment_result['score'] | |
| sentiment_mapping = { | |
| "LABEL_0": "Negative", | |
| "LABEL_1": "Neutral", | |
| "LABEL_2": "Positive" | |
| } | |
| return sentiment_mapping.get(label, "Unknown"), score | |
| except Exception as e: | |
| return f"Error analyzing sentiment: {e}", None | |
| # Function to extract topic using zero-shot classification | |
| def extract_topic(text): | |
| try: | |
| topic_result = topic_pipeline(text, TOPIC_LABELS) | |
| top_topic = topic_result["labels"][0] | |
| confidence = topic_result["scores"][0] | |
| return top_topic, confidence | |
| except Exception as e: | |
| return f"Error extracting topic: {e}", None | |
| # Function to determine if the user's query is about the dataset | |
| def is_dataset_query(text): | |
| keywords = ["dataset", "data", "historical", "csv", "stored"] | |
| text_lower = text.lower() | |
| for keyword in keywords: | |
| if keyword in text_lower: | |
| return True | |
| return False | |
| # Function to retrieve insights from the dataset stored in MongoDB | |
| def get_dataset_insights(): | |
| try: | |
| collection = get_mongo_client() | |
| data = list(collection.find({}, {"_id": 0})) | |
| if not data: | |
| return "The dataset in MongoDB is empty." | |
| df = pd.DataFrame(data) | |
| # Map the sentiment labels from sentiment140.csv: 0 -> Negative, 2 -> Neutral, 4 -> Positive. | |
| sentiment_mapping = {0: "Negative", 2: "Neutral", 4: "Positive"} | |
| if "target" in df.columns: | |
| df['sentiment_label'] = df['target'].apply(lambda x: sentiment_mapping.get(int(x), "Unknown")) | |
| summary = df['sentiment_label'].value_counts().to_dict() | |
| summary_str = ", ".join([f"{k}: {v}" for k, v in summary.items()]) | |
| return f"The dataset sentiment distribution is: {summary_str}." | |
| else: | |
| return "The dataset does not have a 'target' field." | |
| except Exception as e: | |
| return f"Error retrieving dataset insights: {e}" | |
| # Function to generate AI response along with sentiment and topic analysis | |
| def chatbot_response(user_prompt): | |
| if not user_prompt: | |
| return None, None, None, None, None | |
| # Check if the query is about the dataset | |
| if is_dataset_query(user_prompt): | |
| dataset_insights = get_dataset_insights() | |
| sentiment_label, sentiment_confidence = analyze_sentiment(user_prompt) | |
| topic_label, topic_confidence = extract_topic(user_prompt) | |
| return dataset_insights, sentiment_label, sentiment_confidence, topic_label, topic_confidence | |
| else: | |
| try: | |
| # Generate AI response using Gemini | |
| model_gen = genai.GenerativeModel("gemini-1.5-pro") | |
| ai_response = model_gen.generate_content(user_prompt) | |
| sentiment_label, sentiment_confidence = analyze_sentiment(user_prompt) | |
| topic_label, topic_confidence = extract_topic(user_prompt) | |
| return ai_response.text, sentiment_label, sentiment_confidence, topic_label, topic_confidence | |
| except Exception as e: | |
| return f"❌ Error: {e}", None, None, None, None | |