Update app.py

app.py

@@ -1,45 +1,35 @@
 import gradio as gr
 import torch
-from transformers import T5Tokenizer, T5ForConditionalGeneration
+from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
 
-# Load T5 model and tokenizer
-model_name = "t5-base"  # Use a smaller model for faster inference
-tokenizer = T5Tokenizer.from_pretrained(model_name)
-model = T5ForConditionalGeneration.from_pretrained(model_name)
+# Load DistilBERT model and tokenizer
+model_name = "bhadresh-savani/distilbert-base-uncased-finetuned-sentiment"
+tokenizer = DistilBertTokenizer.from_pretrained(model_name)
+model = DistilBertForSequenceClassification.from_pretrained(model_name)
 
 # Use GPU if available
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
-# Grammar correction function
-def correct_grammar(text):
-    input_text = f"correct: {text}"
-    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
-    
-    # Generate corrected text
-    output_ids = model.generate(input_ids, max_length=512, num_beams=5, early_stopping=True)
-    corrected_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    
-    return corrected_text
-
-# Gradio interface function
-def correct_grammar_interface(text):
-    corrected_text = correct_grammar(text)
-    return corrected_text
+# Define the prediction function
+def predict_sentiment(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+    outputs = model(**inputs)
+    predictions = torch.argmax(outputs.logits, dim=-1)
+    return predictions.item()
 
 # Gradio interface
-with gr.Blocks() as grammar_app:
-    gr.Markdown("<h1>Fast Grammar Correction with T5</h1>")
+with gr.Blocks() as sentiment_app:
+    gr.Markdown("<h1>Sentiment Analysis with DistilBERT</h1>")
     
-    with gr.Row():
-        input_box = gr.Textbox(label="Input Text", placeholder="Enter text to be corrected", lines=4)
-        output_box = gr.Textbox(label="Corrected Text", placeholder="Corrected text will appear here", lines=4)
+    input_box = gr.Textbox(label="Input Text", placeholder="Enter text to analyze sentiment")
+    output_box = gr.Textbox(label="Sentiment Result", placeholder="Sentiment result will appear here")
 
-    submit_button = gr.Button("Correct Grammar")
+    submit_button = gr.Button("Analyze Sentiment")
     
     # Button click event
-    submit_button.click(fn=correct_grammar_interface, inputs=input_box, outputs=output_box)
+    submit_button.click(fn=predict_sentiment, inputs=input_box, outputs=output_box)
 
 # Launch the app
 if __name__ == "__main__":
-    grammar_app.launch()
+    sentiment_app.launch()