Updated app.py

app.py

@@ -1,63 +1,53 @@
-from flask import Flask, request, jsonify
+import gradio as gr
 from transformers import PegasusForConditionalGeneration, PegasusTokenizer
 import torch
 import os
 
-# Initialize Flask App
-app = Flask(__name__)
+# --- Global variables to cache the model ---
+# We use a dictionary to hold the model and tokenizer to manage state.
+model_cache = {
+    "model": None,
+    "tokenizer": None
+}
 
-# --- Global variables for the model and tokenizer ---
-# We will load these lazily (on the first request) to avoid startup timeouts.
-model = None
-tokenizer = None
-model_loaded = False
-
-# --- Function to Load the Model ---
+# --- Function to Load the Model (only runs once) ---
 def load_model():
-    """Loads the model and tokenizer into the global variables."""
-    global model, tokenizer, model_loaded
-    
-    # This points to the NEW repository with the corrected model files.
-    model_repo_id = "TheOCEAN/My_Text_Summarizer_Portable" 
-
-    print("--- LAZY LOADING MODEL ---")
-    try:
-        print(f"⬇️ Loading model and tokenizer from '{model_repo_id}'...")
-        tokenizer = PegasusTokenizer.from_pretrained(model_repo_id)
-        model = PegasusForConditionalGeneration.from_pretrained(model_repo_id)
-        
-        # Move model to CPU to ensure it works on free tier hardware
-        model.to(torch.device("cpu"))
-        
-        model_loaded = True
-        print("✅ Model and tokenizer loaded successfully.")
-        print("--- Model Loading Complete ---")
-    except Exception as e:
-        model_loaded = False
-        print(f"❌❌❌ FATAL ERROR DURING MODEL LOADING: {e} ❌❌❌")
-        # This will be visible in the Space logs if something goes wrong.
+    """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) ---")
+        try:
+            # This points to the NEW repository with the corrected model files.
+            model_repo_id = "TheOCEAN/My_Text_Summarizer_Portable" 
+            print(f"⬇️ Loading model and tokenizer from '{model_repo_id}'...")
+            
+            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
+            model_cache["model"] = model
+            model_cache["tokenizer"] = tokenizer
+            
+            print("✅ 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.")
+            
+    return model_cache["model"], model_cache["tokenizer"]
 
-# --- API Route for Summarization ---
-@app.route('/summarize', methods=['POST'])
-def summarize_endpoint():
-    global model, tokenizer, model_loaded
-
-    # If the model isn't loaded yet, load it now.
-    # This happens only on the very first request.
-    if not model_loaded:
-        load_model()
+# --- The core summarization function that Gradio will expose as an 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."
     
-    # If loading failed, return an error.
-    if not model or not tokenizer:
-        return jsonify({'error': 'Model is not available. Check server logs for loading errors.'}), 503
-
     try:
-        json_data = request.get_json()
-        text = json_data.get('text', '')
-        if not text:
-            return jsonify({'error': 'No text provided.'}), 400
-
-        # Use the loaded model and tokenizer to summarize
+        # 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(
@@ -67,23 +57,25 @@ def summarize_endpoint():
             early_stopping=True
         )
         summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-        
-        return jsonify({'summary': summary})
-
+        return summary
     except Exception as e:
-        print(f"An error occurred during summarization: {e}")
-        return jsonify({'error': 'Failed to process the request.'}), 500
+        print(f"Error during summarization: {e}")
+        # Gradio can handle returning error messages to the API caller
+        return f"An error occurred during processing: {e}"
 
-# --- Root Route to Confirm Server is Running ---
-@app.route('/')
-def home():
-    # This route will respond instantly, satisfying the health check.
-    return "Summarizer API is running. Model will be loaded on the first /summarize request."
+# --- 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.
+demo = gr.Interface(
+    fn=summarize_text,
+    inputs=gr.Textbox(lines=15, placeholder="Enter a long text or article to summarize here..."),
+    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."
+)
 
-# --- Main Execution ---
-if __name__ == '__main__':
-    # We no longer load the model on startup.
-    # The port needs to be 7860 for Hugging Face Spaces.
-    port = int(os.environ.get("PORT", 7860))
-    app.run(host='0.0.0.0', port=port)
+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()