Update app.py

app.py

@@ -1,52 +1,70 @@
-# Required: pip install gradio transformers accelerate optimum onnxruntime onnx
+# pip install gradio transformers onnxruntime optimum torch
 
 import gradio as gr
 import torch
 from transformers import AutoTokenizer
-from optimum.onnxruntime import ORTModelForSeq2SeqLM
-from optimum.pipelines import pipeline
+from optimum.onnxruntime import ORTModelForSeq2SeqLM, ORTOptimizer, ORTQuantizer
+from optimum.onnxruntime.configuration import AutoOptimizationConfig
+import onnxruntime as ort
 
-# Load ONNX-optimized model and tokenizer
+# Step 1: Load & optimize the ONNX model
 model_name = "Rahmat82/t5-small-finetuned-summarization-xsum"
-model = ORTModelForSeq2SeqLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+model = ORTModelForSeq2SeqLM.from_pretrained(model_name, export=True)
+
+optimizer = ORTOptimizer.from_pretrained(model)
+opt_config = AutoOptimizationConfig.O2()  # graph fusions and transformer-specific optimizations
+optimizer.optimize(save_dir="optimized_model", optimization_config=opt_config)
+optimized_model = ORTModelForSeq2SeqLM.from_pretrained("optimized_model")
+
+# Step 2: Apply dynamic INT8 quantization for CPU
+quantizer = ORTQuantizer.from_pretrained(optimized_model)
+opt_q = quantizer.quantize(
+    save_dir="quantized_model",
+    quantization_config=AutoOptimizationConfig.O2().quantization_config,  # dynamic quant
+)
+model = ORTModelForSeq2SeqLM.from_pretrained("quantized_model")
 
-# Build a fast summarization pipeline
-device = 0 if torch.cuda.is_available() else -1
+# Step 3: Set up ONNXRuntime Session options for CPU multi-threading
+sess_options = ort.SessionOptions()
+sess_options.intra_op_num_threads = min(4, torch.get_num_threads())  # 4 threads for inference
+sess_options.inter_op_num_threads = 1
+sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+# Rebuild pipeline with optimized quantized model on CPU
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+summarizer = gradio_pipeline = None
+from optimum.pipelines import pipeline
 summarizer = pipeline(
     task="summarization",
     model=model,
     tokenizer=tokenizer,
-    device=device,
-    batch_size=16,  # increased batch size for higher throughput
+    framework="pt",
+    ort_session_options=sess_options,
+    device=-1,
+    batch_size=8,
 )
 
-# Speed-optimized summarization function
 def summarize_text(text):
     text = text.strip()
     if not text:
         return "Please enter some text."
-    
-    # Encode with truncation (max_length=1024)
     inputs = tokenizer.encode(text, max_length=1024, truncation=True, return_tensors="pt")
-    decoded_input = tokenizer.decode(inputs[0], skip_special_tokens=True)
-
-    # Generate summary with tighter bounds
+    decoded = tokenizer.decode(inputs[0], skip_special_tokens=True)
     summary = summarizer(
-        decoded_input,
-        min_length=69,   # lower min length for faster generation
+        decoded,
+        min_length=60,
         max_length=120,
         do_sample=False
     )
     return summary[0]["summary_text"]
 
-# Gradio interface
+# Gradio UI
 app = gr.Interface(
     fn=summarize_text,
-    inputs=gr.Textbox(lines=12, placeholder="Paste long text here...", label="Input Text"),
+    inputs=gr.Textbox(lines=12, label="Input Text"),
     outputs=gr.Textbox(label="Summary"),
-    title="⚡ Fast ONNX T5 Summarizer",
-    description="ONNX-accelerated T5-small model for quick, medium-length summarization (up to 1,024 tokens)."
+    title="⚙️ CPU-Optimized ONNX T5 Summarizer",
+    description="Uses graph optimizations, INT8 quantization, and threading tweaks for fast CPU performance."
 )
 
 if __name__ == "__main__":