Update app.py

app.py

@@ -1,25 +1,213 @@
-    model = keras.Sequential([
-        vectorizer,
-        Embedding(input_dim=vocab_size, output_dim=64),
-        GlobalAveragePooling1D(),
-        Dense(64, activation="relu"),
-        Dense(len(df['label'].unique()), activation="softmax")  # Adjust for number of classes
-    ])
-    
-    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
-    return model
-# Train the model
-model = build_model()
-model.fit(train_ds.batch(32), epochs=10, validation_data=test_ds.batch(32))
-# Function to make predictions
-def predict(text):
+import gradio as gr
+import pandas as pd
+import numpy as np
+import torch
+from datasets import load_dataset
+from sklearn.model_selection import train_test_split
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import accuracy_score, classification_report
+from sklearn.preprocessing import LabelEncoder
+import matplotlib.pyplot as plt
+from imblearn.over_sampling import SMOTE
+import plotly.express as px
+import plotly.graph_objects as go
+import warnings
+
+# Suppress warnings
+warnings.filterwarnings("ignore")
+
+# Load dataset
+print("Loading dataset...")
+ds = load_dataset("uhoui/text-tone-classifier")
+
+# Convert to pandas DataFrame
+df = pd.DataFrame(ds["train"])
+
+# Print dataset statistics
+print(f"Dataset size: {len(df)} entries")
+print(f"Columns: {df.columns}")
+
+# Check class distribution
+label_counts = df['label'].value_counts()
+print("\nClass distribution:")
+print(label_counts)
+
+# Encode labels
+label_encoder = LabelEncoder()
+df['label_encoded'] = label_encoder.fit_transform(df['label'])
+num_classes = len(label_encoder.classes_)
+
+# Split the data
+X_train, X_test, y_train, y_test = train_test_split(
+    df['text'], 
+    df['label_encoded'], 
+    test_size=0.2, 
+    random_state=42, 
+    stratify=df['label_encoded'] if len(df) > 10 else None  # Only stratify if we have enough samples
+)
+
+# Feature extraction using TF-IDF
+print("Creating TF-IDF features...")
+tfidf = TfidfVectorizer(max_features=5000)
+X_train_tfidf = tfidf.fit_transform(X_train)
+X_test_tfidf = tfidf.transform(X_test)
+
+# Handle class imbalance using SMOTE
+print("Applying SMOTE to handle class imbalance...")
+try:
+    smote = SMOTE(random_state=42)
+    X_train_resampled, y_train_resampled = smote.fit_resample(X_train_tfidf, y_train)
+    print(f"After SMOTE: {X_train_resampled.shape}")
+except ValueError as e:
+    print(f"SMOTE error: {e}. Using original data.")
+    X_train_resampled, y_train_resampled = X_train_tfidf, y_train
+
+# Train a logistic regression model
+print("Training model...")
+model = LogisticRegression(C=10, max_iter=1000, n_jobs=-1, solver='lbfgs', multi_class='multinomial')
+model.fit(X_train_resampled, y_train_resampled)
+
+# Evaluate model
+y_pred = model.predict(X_test_tfidf)
+accuracy = accuracy_score(y_test, y_pred)
+print(f"Model accuracy: {accuracy:.4f}")
+
+# Function to predict tone with probabilities
+def predict_tone(text):
     # Vectorize the input text
-    vectorized_text = vectorizer([text])  # Use the vectorizer to transform the input
-    prediction = model.predict(vectorized_text)  # Pass the vectorized input to the model
-    predicted_label = tf.argmax(prediction, axis=1).numpy()[0]
-    return df['label'].cat.categories[predicted_label]
+    text_tfidf = tfidf.transform([text])
+    
+    # Get prediction probabilities
+    probs = model.predict_proba(text_tfidf)[0]
+    
+    # Get the predicted class and its probability
+    pred_class_idx = np.argmax(probs)
+    pred_class = label_encoder.inverse_transform([pred_class_idx])[0]
+    
+    # Create results dictionary with all probabilities
+    results = {}
+    for i, label in enumerate(label_encoder.classes_):
+        results[label] = float(probs[i])
+    
+    # Sort results by probability (descending)
+    sorted_results = {k: v for k, v in sorted(results.items(), key=lambda item: item[1], reverse=True)}
+    
+    # Create visualization
+    top_n = 5  # Show top 5 emotions
+    top_labels = list(sorted_results.keys())[:top_n]
+    top_probs = list(sorted_results.values())[:top_n]
+    
+    # Generate colors based on probability (higher probability = more intense color)
+    colors = ["rgba(64, 128, 255, " + str(min(1.0, p + 0.3)) + ")" for p in top_probs]
+    
+    fig = go.Figure()
+    fig.add_trace(go.Bar(
+        x=top_probs,
+        y=top_labels,
+        orientation='h',
+        marker_color=colors,
+        text=[f"{p:.1%}" for p in top_probs],
+        textposition='auto'
+    ))
+    
+    fig.update_layout(
+        title="Emotion Probability",
+        xaxis_title="Probability",
+        yaxis_title="Emotion",
+        height=400,
+        margin=dict(l=20, r=20, t=40, b=20),
+        xaxis=dict(range=[0, 1])
+    )
+    
+    # Get example texts for the predicted emotion
+    example_texts = df[df['label'] == pred_class]['text'].sample(min(3, len(df[df['label'] == pred_class]))).tolist()
+    
+    return pred_class, sorted_results, fig, example_texts
+
+# Function to handle the example display
+def get_tone_examples(tone):
+    examples = df[df['label'] == tone]['text'].sample(min(5, len(df[df['label'] == tone]))).tolist()
+    return examples
+
 # Gradio interface
-iface = gr.Interface(fn=predict, inputs="text", outputs="text", title="Text Tone Sentiment Analysis",
-                     description="Enter a text to analyze its tone (e.g., joy, depression, contentment).")
+def analyze_tone(text, selected_tone=None):
+    if not text:
+        return "Please enter some text to analyze.", {}, None, []
+    
+    # If a tone is selected from the dropdown, show examples
+    if selected_tone and not text:
+        examples = get_tone_examples(selected_tone)
+        return f"Examples of '{selected_tone}' tone:", {}, None, examples
+    
+    # Otherwise, analyze the text
+    predicted_tone, all_probs, fig, examples = predict_tone(text)
+    
+    # Format the result message
+    message = f"The predicted tone is: **{predicted_tone}**"
+    
+    return message, all_probs, fig, examples
+
+# Create the Gradio interface
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
+    gr.Markdown("# Text Tone Analyzer")
+    gr.Markdown("Enter text to analyze its emotional tone.")
+    
+    with gr.Row():
+        with gr.Column(scale=3):
+            text_input = gr.Textbox(
+                label="Enter your text here",
+                placeholder="Type something to analyze its emotional tone...",
+                lines=5
+            )
+            analyze_button = gr.Button("Analyze Tone", variant="primary")
+            
+        with gr.Column(scale=2):
+            # Dropdown to select example tones
+            tone_dropdown = gr.Dropdown(
+                choices=sorted(df['label'].unique().tolist()),
+                label="Or select a tone to see examples"
+            )
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            result_message = gr.Markdown()
+        
+    with gr.Row():
+        with gr.Column(scale=2):
+            plot_output = gr.Plot(label="Tone Probabilities")
+        with gr.Column(scale=1):
+            all_probs_output = gr.JSON(label="All Probabilities")
+            
+    with gr.Row():
+        examples_output = gr.Dataframe(
+            headers=["Examples of similar texts"],
+            datatype=["str"],
+            label="Example texts with similar tone"
+        )
+    
+    # Set up event handlers
+    analyze_button.click(
+        fn=analyze_tone,
+        inputs=[text_input, None],
+        outputs=[result_message, all_probs_output, plot_output, examples_output]
+    )
+    
+    tone_dropdown.change(
+        fn=get_tone_examples,
+        inputs=[tone_dropdown],
+        outputs=[examples_output]
+    )
+    
+    # Add example inputs
+    examples = [
+        ["I'm so excited about this new project!"],
+        ["I'm feeling quite down today and nothing seems to work."],
+        ["The movie was interesting, but I'm not sure if I liked it."],
+        ["I can't believe what just happened! This is outrageous!"]
+    ]
+    gr.Examples(examples=examples, inputs=text_input)
+
+# Launch the app
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()