Update app.py

app.py

@@ -1,35 +1,65 @@
-import streamlit as st
-from transformers import DistilBertTokenizer, TFDistilBertForSequenceClassification
-import tensorflow as tf
-
-# Load the pre-trained model and tokenizer using the correct Hugging Face model repo ID
-model_path = 'shukdevdatta123/Dreaddit_DistillBert_Stress_Model'
-loaded_model = TFDistilBertForSequenceClassification.from_pretrained(model_path)
-loaded_tokenizer = DistilBertTokenizer.from_pretrained(model_path)
-
-# Define the prediction function
-def predict_with_loaded_model(in_sentences):
-    labels = ["non-stress", "stress"]
-    inputs = loaded_tokenizer(in_sentences, return_tensors="tf", padding=True, truncation=True, max_length=512)
-    predictions = loaded_model(inputs)
-    predicted_labels = tf.argmax(predictions.logits, axis=-1).numpy()
-    predicted_probs = tf.nn.softmax(predictions.logits, axis=-1).numpy()
-
-    return [{"text": sentence, "confidence": probs.tolist(), "label": labels[label]} for sentence, label, probs in zip(in_sentences, predicted_labels, predicted_probs)]
-
-# Streamlit interface
-st.title("Stress Prediction with DistilBERT")
-
-# Add a text input box for the user to enter a sentence
-user_input = st.text_area("Enter a sentence or text:", "")
-
-# When the user clicks "Predict", run the prediction function
-if st.button("Predict"):
-    if user_input:
-        # Make the prediction using the model
-        prediction = predict_with_loaded_model([user_input])[0]
-        st.write(f"Text: {prediction['text']}")
-        st.write(f"Prediction: {prediction['label']}")
-        # st.write(f"Confidence: {prediction['confidence']}")
-    else:
-        st.write("Please enter a sentence to predict.")
+import gradio as gr
+from transformers import DistilBertTokenizer, TFDistilBertForSequenceClassification
+import tensorflow as tf
+
+# Load the pre-trained model and tokenizer using the correct Hugging Face model repo ID
+model_path = 'shukdevdatta123/Dreaddit_DistillBert_Stress_Model'
+loaded_model = TFDistilBertForSequenceClassification.from_pretrained(model_path)
+loaded_tokenizer = DistilBertTokenizer.from_pretrained(model_path)
+
+# Define the prediction function
+def predict_with_loaded_model(in_sentence):
+    if not in_sentence:
+        return "Please enter a sentence to predict.", None
+    
+    labels = ["non-stress", "stress"]
+    inputs = loaded_tokenizer([in_sentence], return_tensors="tf", padding=True, truncation=True, max_length=512)
+    predictions = loaded_model(inputs)
+    predicted_labels = tf.argmax(predictions.logits, axis=-1).numpy()
+    predicted_probs = tf.nn.softmax(predictions.logits, axis=-1).numpy()
+    
+    label = labels[predicted_labels[0]]
+    confidence = predicted_probs[0].tolist()
+    
+    # Create a dictionary for the label output
+    result = {labels[i]: confidence[i] for i in range(len(labels))}
+    
+    return f"Prediction: {label}", result
+
+# Create Gradio interface
+with gr.Blocks(title="Stress Prediction with DistilBERT") as demo:
+    gr.Markdown("# Stress Prediction with DistilBERT")
+    gr.Markdown("Enter a sentence or text to predict whether it indicates stress or non-stress.")
+    
+    with gr.Row():
+        with gr.Column():
+            input_text = gr.Textbox(
+                label="Enter a sentence or text:",
+                placeholder="Type your text here...",
+                lines=5
+            )
+            predict_btn = gr.Button("Predict", variant="primary")
+        
+        with gr.Column():
+            output_text = gr.Textbox(label="Result", lines=2)
+            output_confidence = gr.Label(label="Confidence Scores", num_top_classes=2)
+    
+    predict_btn.click(
+        fn=predict_with_loaded_model,
+        inputs=input_text,
+        outputs=[output_text, output_confidence]
+    )
+    
+    # Add examples
+    gr.Examples(
+        examples=[
+            ["I feel overwhelmed with all the work I have to do."],
+            ["I had a great day at the park with my friends."],
+            ["The deadline is tomorrow and I haven't even started."]
+        ],
+        inputs=input_text
+    )
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch()