Update app.py

app.py

@@ -1,44 +1,50 @@
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline, BitsAndBytesConfig
 import gradio as gr
-
-model_name = "sshleifer/distilbart-cnn-12-6"
-
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# Setup 8-bit quantization
-bnb_config = BitsAndBytesConfig(
-    load_in_8bit=True
-)
-
-# Load model
-model = AutoModelForSeq2SeqLM.from_pretrained(
-    model_name,
-    quantization_config=bnb_config,
-    device_map="auto"  # automatically maps to CPU/GPU
-)
-
-# Create pipeline WITHOUT device argument
-summarizer = pipeline(
-    "summarization",
-    model=model,
-    tokenizer=tokenizer
-)
-
-# 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_ui,
-    inputs=[gr.Textbox(label="Input text to summarize", lines=6)],
-    outputs=[gr.Textbox(label="Summarized text", lines=4)],
-    title="8-bit CPU Text Summarizer",
-    description="Summarize your text fast on CPU using 8-bit quantization"
+from optimum.onnxruntime import ORTModelForSeq2SeqLM
+from transformers import AutoTokenizer, pipeline
+
+# Load ONNX model
+def create_fast_summarizer():
+    model = ORTModelForSeq2SeqLM.from_pretrained(
+        "onnx-community/bart-large-cnn-ONNX",
+        encoder_file_name="encoder_model_q4.onnx",
+        decoder_file_name="decoder_model_q4.onnx",
+        provider="CPUExecutionProvider",
+        use_io_binding=True
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        "onnx-community/bart-large-cnn-ONNX",
+        use_fast=True
+    )
+    return pipeline(
+        "summarization",
+        model=model,
+        tokenizer=tokenizer,
+        device=-1
+    )
+
+summarizer = create_fast_summarizer()
+
+# Summarize function with prompt + tuned params
+def summarize_text(text):
+    prompt = "Summarize the key events, including casualties and political context:\n" + text
+    result = summarizer(
+        prompt,
+        max_length=160,
+        min_length=80,
+        do_sample=False,
+        num_beams=6,
+        length_penalty=1.5,
+        early_stopping=True
+    )
+    return result[0]['summary_text']
+
+# Build Gradio interface
+app = gr.Interface(
+    fn=summarize_text,
+    inputs=gr.Textbox(lines=15, placeholder="Paste your text here..."),
+    outputs="text",
+    title="ONNX Summarizer 🚀",
+    description="Paste any news or article text and get a concise, context-rich summary."
 )
 
-demo.launch()
+app.launch()