Spaces:
Sleeping
Sleeping
| 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() | |