Update app.py

app.py

@@ -1,35 +1,20 @@
-import gradio as gr
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 import torch
-from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
 
-# 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)
+# Load the T5 tokenizer and model
+model_name = "t5-small"  # You can use any T5 model available
+tokenizer = T5Tokenizer.from_pretrained(model_name)
+model = T5ForConditionalGeneration.from_pretrained(model_name)
 
-# Use GPU if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+# Example function to use the model
+def summarize(text):
+    # Tokenize the input text
+    inputs = tokenizer.encode("summarize: " + text, return_tensors="pt", max_length=512, truncation=True)
+    # Generate summary
+    outputs = model.generate(inputs, max_length=150, min_length=30, length_penalty=2.0, num_beams=4, early_stopping=True)
+    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return summary
 
-# 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 sentiment_app:
-    gr.Markdown("<h1>Sentiment Analysis with DistilBERT</h1>")
-    
-    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("Analyze Sentiment")
-    
-    # Button click event
-    submit_button.click(fn=predict_sentiment, inputs=input_box, outputs=output_box)
-
-# Launch the app
-if __name__ == "__main__":
-    sentiment_app.launch()
+# Example usage
+text_to_summarize = "Your input text goes here."
+print(summarize(text_to_summarize))