Updated app.py

app.py

@@ -1,10 +1,10 @@
 import gradio as gr
 from transformers import PegasusForConditionalGeneration, PegasusTokenizer
 import torch
+from huggingface_hub import hf_hub_download
 import os
 
-# --- Global variables to cache the model ---
-# We use a dictionary to hold the model and tokenizer to manage state.
+# --- Global cache for the model ---
 model_cache = {
     "model": None,
     "tokenizer": None
@@ -12,42 +12,50 @@ model_cache = {
 
 # --- Function to Load the Model (only runs once) ---
 def load_model():
-    """Loads the model and tokenizer if they are not already loaded."""
-    # This check prevents reloading on every request
-    if model_cache["model"] is None or model_cache["tokenizer"] is None:
-        print("--- LAZY LOADING MODEL (First Request) ---")
+    """Loads the quantized model and tokenizer if they are not already loaded."""
+    if model_cache["model"] is None:
+        print("--- LAZY LOADING QUANTIZED MODEL (First Request) ---")
         try:
-            # This points to the NEW repository with the corrected model files.
-            model_repo_id = "TheOCEAN/My_Text_Summarizer_Portable" 
+            # This now points to the final, smaller, quantized model repository
+            model_repo_id = "TheOCEAN/My_Text_Summarizer_Quantized" 
             print(f"⬇️ Loading model and tokenizer from '{model_repo_id}'...")
             
+            # 1. Load the tokenizer
             tokenizer = PegasusTokenizer.from_pretrained(model_repo_id)
-            model = PegasusForConditionalGeneration.from_pretrained(model_repo_id)
-            model.to(torch.device("cpu")) # Ensure it runs on CPU
             
-            # Store the loaded model and tokenizer in our cache
+            # 2. Create the base model structure
+            model = PegasusForConditionalGeneration.from_pretrained("google/pegasus-xsum")
+            
+            # 3. Apply the same quantization structure to the base model
+            model = torch.quantization.quantize_dynamic(
+                model, {torch.nn.Linear}, dtype=torch.qint8
+            )
+            
+            # 4. Download and load your fine-tuned quantized weights
+            weights_path = hf_hub_download(repo_id=model_repo_id, filename="quantized_weights.pth")
+            model.load_state_dict(torch.load(weights_path, map_location="cpu"))
+            
+            model.to(torch.device("cpu"))
+            
             model_cache["model"] = model
             model_cache["tokenizer"] = tokenizer
             
-            print("✅ Model and tokenizer loaded successfully.")
+            print("✅ Quantized model and tokenizer loaded successfully.")
         except Exception as e:
             print(f"❌❌❌ FATAL ERROR DURING MODEL LOADING: {e} ❌❌❌")
-            # Raise a Gradio-specific error that will be shown in the UI
-            raise gr.Error("Model failed to load. Please check the Space logs for details.")
+            raise gr.Error("Model failed to load. Please check the Space logs.")
             
     return model_cache["model"], model_cache["tokenizer"]
 
-# --- The core summarization function that Gradio will expose as an API ---
+# --- The core summarization function for the API ---
 def summarize_text(text):
     """Takes text input and returns a summary."""
     if not text or not text.strip():
         return "Please provide some text to summarize."
     
     try:
-        # Load the model (or get it from the cache if already loaded)
         model, tokenizer = load_model()
         
-        # Tokenize and generate the summary
         inputs = tokenizer(text, max_length=1024, truncation=True, return_tensors="pt")
         
         summary_ids = model.generate(
@@ -60,22 +68,19 @@ def summarize_text(text):
         return summary
     except Exception as e:
         print(f"Error during summarization: {e}")
-        # Gradio can handle returning error messages to the API caller
         return f"An error occurred during processing: {e}"
 
 # --- Create and launch the Gradio Interface ---
-# This creates a simple UI on the Space page for testing, 
-# and automatically creates a background API endpoint that your local app will call.
+# The api_name allows us to call this function like a REST API
 demo = gr.Interface(
     fn=summarize_text,
-    inputs=gr.Textbox(lines=15, placeholder="Enter a long text or article to summarize here..."),
+    inputs=gr.Textbox(lines=15, placeholder="Enter text to summarize..."),
     outputs="text",
-    title="AI Text Summarizer",
-    description="This is a demo of a fine-tuned Pegasus model. The model is loaded on the first request, which may take a few minutes. Subsequent requests will be fast."
+    title="AI Text Summarizer (Quantized)",
+    description="This demo uses a smaller, faster, quantized version of the fine-tuned Pegasus model.",
+    api_name="summarize" 
 )
 
 if __name__ == "__main__":
-    # The launch() command tells Gradio to start the web server.
-    # It will automatically use the correct port for Hugging Face Spaces.
     demo.launch()