Update app.py

app.py

@@ -1,36 +1,41 @@
-import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
 import gradio as gr
-from transformers import pipeline
+import torch
 
-# CPU device
-device = -1
+model_name = "sshleifer/distilbart-cnn-12-6"
 
-# Optimized pipeline
-text_summary = pipeline(
-    "summarization",
-    model="sshleifer/distilbart-cnn-12-6",
-    device=device,
-    torch_dtype=torch.float32  # bfloat16 not always supported on CPU
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Load model in 8-bit (CPU-friendly)
+model = AutoModelForSeq2SeqLM.from_pretrained(
+    model_name,
+    load_in_8bit=True,  # reduces memory & speeds up CPU
+    device_map="auto"   # automatically maps to CPU
 )
 
-# Optional: compile model if PyTorch >=2.1
-try:
-    text_summary.model = torch.compile(text_summary.model)
-except:
-    pass
+# Create pipeline
+summarizer = pipeline(
+    "summarization",
+    model=model,
+    tokenizer=tokenizer,
+    device=-1
+)
 
-def summary(input_text):
-    with torch.no_grad():
-        output = text_summary(input_text, max_length=512, min_length=30, truncation=True)
+# Function for Gradio
+def summary_ui(text):
+    output = summarizer(text, max_length=512, min_length=30, truncation=True)
     return output[0]['summary_text']
 
 gr.close_all()
 
+# Gradio interface
 demo = gr.Interface(
-    fn=summary,
+    fn=summary_ui,
     inputs=[gr.Textbox(label="Input text to summarize", lines=6)],
     outputs=[gr.Textbox(label="Summarized text", lines=4)],
-    title="@GenAILearniverse Project 1: Text Summarizer",
-    description="THIS APPLICATION WILL BE USED TO SUMMARIZE THE TEXT"
+    title="8-bit CPU Text Summarizer",
+    description="Summarize your text fast on CPU using 8-bit quantization"
 )
+
 demo.launch()