import gradio as gr
from transformers import PegasusForConditionalGeneration, PegasusTokenizer
import torch
from huggingface_hub import hf_hub_download
import os

# --- Global cache for the model ---
model_cache = {
    "model": None,
    "tokenizer": None
}

# --- Function to Load the Model (only runs once) ---
def load_model():
    """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 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)
            
            # 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("✅ Quantized model and tokenizer loaded successfully.")
        except Exception as e:
            print(f"❌❌❌ FATAL ERROR DURING MODEL LOADING: {e} ❌❌❌")
            raise gr.Error("Model failed to load. Please check the Space logs.")
            
    return model_cache["model"], model_cache["tokenizer"]

# --- 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:
        model, tokenizer = load_model()
        
        inputs = tokenizer(text, max_length=1024, truncation=True, return_tensors="pt")
        
        summary_ids = model.generate(
            inputs.input_ids, 
            num_beams=4, 
            max_length=150, 
            early_stopping=True
        )
        summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
        return summary
    except Exception as e:
        print(f"Error during summarization: {e}")
        return f"An error occurred during processing: {e}"

# --- Create and launch the Gradio Interface ---
# 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 text to summarize..."),
    outputs="text",
    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__":
    demo.launch()