google

app.py

@@ -3,14 +3,13 @@ import re
 import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from nltk.corpus import stopwords
-from spaces import GPU
 import nltk
 
-# Download stopwords if not already
+# Download NLTK stopwords
 nltk.download('stopwords')
 stop_words = set(stopwords.words('english'))
 
-# Model choices
+# Best lightweight summarization models
 model_choices = {
     "DistilBART CNN (sshleifer/distilbart-cnn-12-6)": "sshleifer/distilbart-cnn-12-6",
     "T5 Small (t5-small)": "t5-small",
@@ -21,12 +20,27 @@ model_choices = {
 
 model_cache = {}
 
+# Clean input text (remove stopwords and SKUs/product codes)
 def clean_text(input_text):
-    cleaned_text = re.sub(r'[^A-Za-z0-9\s]', ' ', input_text)
+    # Remove simple SKU codes (e.g., ST1642, AB1234, etc.)
+    cleaned_text = re.sub(r'\b[A-Za-z]{2,}[0-9]{3,}\b', '', input_text)  # Alphanumeric SKU
+    
+    # Replace special characters with a space
+    cleaned_text = re.sub(r'[^A-Za-z0-9\s]', ' ', cleaned_text)
+    
+    # Tokenize the input text and remove stop words
     words = cleaned_text.split()
     words = [word for word in words if word.lower() not in stop_words]
-    return " ".join(words).strip()
+    
+    # Rebuild the cleaned text
+    cleaned_text = " ".join(words)
+    
+    # Strip leading and trailing spaces
+    cleaned_text = cleaned_text.strip()
+    
+    return cleaned_text
 
+# Load model and tokenizer
 def load_model(model_name):
     if model_name not in model_cache:
         tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -35,16 +49,16 @@ def load_model(model_name):
             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
         )
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model = model.to(device)
-        model_cache[model_name] = (tokenizer, model)
+        model.to(device)
 
-        # Warm up with dummy input
-        dummy_input = tokenizer("summarize: hello world", return_tensors="pt").input_ids.to(device)
+        # Warm-up
+        dummy_input = tokenizer("summarize: warm up", return_tensors="pt").input_ids.to(device)
         model.generate(dummy_input, max_length=10)
 
+        model_cache[model_name] = (tokenizer, model)
     return model_cache[model_name]
 
-@GPU  # 👈 Required for ZeroGPU to allocate GPU when this is called
+# Summarize the text using a selected model
 def summarize_text(input_text, model_label, char_limit):
     if not input_text.strip():
         return "Please enter some text."
@@ -62,8 +76,8 @@ def summarize_text(input_text, model_label, char_limit):
 
     summary_ids = model.generate(
         input_ids,
-        max_length=20,
-        min_length=5,
+        max_length=30,  # Ensure max_length is greater than min_length
+        min_length=5,   # Ensure min_length is less than max_length
         do_sample=False
     )
 
@@ -75,12 +89,12 @@ iface = gr.Interface(
     fn=summarize_text,
     inputs=[
         gr.Textbox(lines=6, label="Enter text to summarize"),
-        gr.Dropdown(choices=list(model_choices.keys()), label="Choose summarization model", value="Pegasus (google/pegasus-xsum)"),
-        gr.Slider(minimum=10, maximum=100, value=65, step=1, label="Max Character Limit")
+        gr.Dropdown(choices=list(model_choices.keys()), label="Choose summarization model", value="DistilBART CNN (sshleifer/distilbart-cnn-12-6)"),
+        gr.Slider(minimum=30, maximum=200, value=80, step=1, label="Max Character Limit")
     ],
     outputs=gr.Textbox(lines=3, label="Summary (truncated to character limit)"),
-    title="Multi-Model Text Summarizer (GPU Ready)",
-    description="Summarize long or short texts using state-of-the-art Hugging Face models with GPU acceleration (ZeroGPU-compatible)."
+    title="🚀 Fast Lightweight Summarizer (GPU Optimized)",
+    description="Summarize text quickly using compact models ideal for low-latency and ZeroGPU Spaces."
 )
 
 iface.launch()