Update app.py

app.py

@@ -1,28 +1,38 @@
+# Required: pip install gradio transformers accelerate optimum onnxruntime onnx
+
 import gradio as gr
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
+from transformers import AutoTokenizer
+from optimum.onnxruntime import ORTModelForSeq2SeqLM
+from optimum.pipelines import pipeline
 
-# Load the model and tokenizer
-model_id = "Rahmat82/t5-small-finetuned-summarization-xsum"
-model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
+# Load ONNX model and tokenizer
+model_name = "Rahmat82/t5-small-finetuned-summarization-xsum"
+model = ORTModelForSeq2SeqLM.from_pretrained(model_name, export=True)
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
 
-# Create the summarization pipeline
-summarizer = pipeline("summarization", model=model, tokenizer=tokenizer)
+# Create summarizer pipeline with Optimum
+summarizer = pipeline("summarization", model=model, tokenizer=tokenizer, device_map="auto", batch_size=12)
 
-# Define the summarization function
+# Define summarization function with 1024 token cap
 def summarize_text(text):
-    if len(text.strip()) == 0:
+    if not text.strip():
         return "Please enter some text."
-    summary = summarizer(text)[0]["summary_text"]
-    return summary
+    
+    # Tokenize and truncate to max 1024 tokens
+    inputs = tokenizer(text, return_tensors="pt", max_length=1024, truncation=True)
+    input_text = tokenizer.decode(inputs["input_ids"][0], skip_special_tokens=True)
+
+    # Summarize the truncated text
+    result = summarizer(input_text)
+    return result[0]["summary_text"]
 
-# Gradio Interface
+# Gradio app
 app = gr.Interface(
     fn=summarize_text,
     inputs=gr.Textbox(lines=15, placeholder="Paste your text here...", label="Input Text"),
     outputs=gr.Textbox(label="Summary"),
-    title="Custom T5 Summarizer 📝",
-    description="Summarize long texts using a fine-tuned T5-small model from Hugging Face"
+    title="🚀 ONNX-Powered T5 Summarizer (1024 tokens)",
+    description="Summarize long text using a fine-tuned ONNX-accelerated T5-small model (max input: 1024 tokens)"
 )
 
 # Launch the app