Update app.py

app.py

@@ -1,22 +1,45 @@
-import streamlit as st
-from transformers import pipeline
+import torch
+import gradio as gr
+from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
 
-# Load the summariztion model pipeline
-summarizer_ntg = pipeline("text2text-generation", model="mrm8488/t5-base-finetuned-summarize-news")
-classifier = pipeline("text-classification", model='Lauraayu/News_Classi_Model', return_all_scores=True)
+# Step 1: Define Summarization Pipeline
+summarizer_ntg = pipeline("summarization", model="mrm8488/t5-base-finetuned-summarize-news")
 
-# Streamlit application title
-st.title("News Classification")
-st.write("Classification for different News types")
+# Step 2: Define Classification Pipeline
+tokenizer_bb = AutoTokenizer.from_pretrained("Lauraayu/News_Classi_Model")
+model_bb = AutoModelForSequenceClassification.from_pretrained("Lauraayu/News_Classi_Model")
 
-# Text input for user to enter the text to classify
-text = st.text_area("Enter the News to classify","")
-
-# Perform text classification when the user clicks the "Classify" button
-if st.button("Classify"):
+def summarize_and_classify(text):
+    # Summarize the article
+    summary = summarizer_ntg(text)[0]['summary_text']
+    
+    # Tokenize the summarized text
+    inputs = tokenizer_bb(summary, return_tensors="pt", truncation=True, padding=True, max_length=512)
+    
+    # Move inputs and model to the same device (GPU or CPU)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    model_bb.to(device)
     
-    # Perform text classification on the input text
-    result0 = summarizer_ntg(text)
-    result = classifier(result0)
+    # Perform classification
+    with torch.no_grad():
+        outputs = model_bb(**inputs)
+    
+    # Get the predicted label
+    predicted_label_id = torch.argmax(outputs.logits, dim=-1).item()
+    label_mapping = model_bb.config.id2label
+    predicted_label = label_mapping[predicted_label_id]
+    
+    return summary, predicted_label
+
+# Create Gradio Interface
+iface = gr.Interface(
+    fn=summarize_and_classify,
+    inputs=gr.inputs.Textbox(lines=10, placeholder="Enter news article text here..."),
+    outputs=[gr.outputs.Textbox(label="Summary"), gr.outputs.Textbox(label="Category")],
+    title="News Article Summarizer and Classifier",
+    description="Enter a news article text and get its summary and category."
+)
 
-    st.write(result)
+# Launch the interface
+iface.launch()