| import gradio as gr |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| MODEL_ID = "arinbalyan/summarization-lora" |
| MAX_LENGTH = 512 |
| MAX_NEW_TOKENS = 150 |
|
|
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| print(f"Loading model from {MODEL_ID}...") |
| model = AutoModelForCausalLM.from_pretrained( |
| MODEL_ID, |
| torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, |
| device_map="auto", |
| trust_remote_code=True, |
| ) |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) |
| if tokenizer.pad_token is None: |
| tokenizer.pad_token = tokenizer.eos_token |
| print("Model loaded successfully.") |
|
|
|
|
| INSTRUCTION_TEMPLATE = "Summarize the following article:\n\n{article}\n\nSummary:" |
|
|
|
|
| def summarize(article_text, temperature=0.3, max_new_tokens=120): |
| """Generate a summary for an article.""" |
| if not article_text or article_text.strip() == "": |
| return "Please enter an article to summarize." |
|
|
| prompt = INSTRUCTION_TEMPLATE.format(article=article_text.strip()) |
|
|
| inputs = tokenizer( |
| prompt, return_tensors="pt", truncation=True, max_length=MAX_LENGTH |
| ) |
| inputs = {k: v.to(device) for k, v in inputs.items()} |
|
|
| with torch.no_grad(): |
| outputs = model.generate( |
| **inputs, |
| max_new_tokens=max_new_tokens, |
| temperature=temperature, |
| top_p=0.9, |
| do_sample=True, |
| pad_token_id=tokenizer.pad_token_id, |
| eos_token_id=tokenizer.eos_token_id, |
| ) |
|
|
| generated = tokenizer.decode(outputs[0], skip_special_tokens=True) |
| if "Summary:" in generated: |
| summary = generated.split("Summary:")[-1].strip() |
| else: |
| summary = generated[len(prompt) :].strip() |
| return summary |
|
|
|
|
| |
| SAMPLE_ARTICLES = [ |
| ( |
| "Technology", |
| "Apple today announced the new MacBook Pro featuring the M4 chip family, " |
| "delivering up to 2x faster performance than the previous generation. " |
| "The new lineup includes 14-inch and 16-inch models with Thunderbolt 5, " |
| "up to 24 hours of battery life, a 12MP Center Stage camera, and a " |
| "stunning Liquid Retina XDR display. Pre-orders begin today with " |
| "availability starting next Friday.", |
| ), |
| ( |
| "Science", |
| "A team of researchers at MIT has developed a new type of battery that " |
| "could revolutionize energy storage for electric vehicles. The solid-state " |
| "battery uses a novel electrolyte material that is both safer and more " |
| "energy-dense than current lithium-ion batteries. In tests, the new battery " |
| "achieved 500 miles of range on a single charge and charged to 80% in just " |
| "15 minutes. The researchers say the technology could be commercially " |
| "available within three years.", |
| ), |
| ( |
| "Environment", |
| "A landmark climate agreement was reached at the COP30 summit in Brazil " |
| "today, with 195 countries committing to reduce methane emissions by 45% " |
| "by 2035. The agreement includes $100 billion in annual funding for " |
| "developing nations to transition to renewable energy. Environmental groups " |
| "hailed the deal as historic but warned that enforcement mechanisms remain " |
| "weak. Critics point out that several major emitters have yet to sign.", |
| ), |
| ] |
|
|
|
|
| with gr.Blocks( |
| title="Text Summarization — SmolLM2 LoRA", |
| theme=gr.themes.Soft(), |
| css=""" |
| footer { display: none !important; } |
| .gradio-container { max-width: 900px; margin: auto; } |
| """, |
| ) as demo: |
| gr.Markdown( |
| """ |
| # 📝 Text Summarization with SmolLM2-1.7B (LoRA Fine-Tuned) |
| |
| Enter an article below to generate a concise summary using a LoRA fine-tuned SmolLM2-1.7B model |
| on the CNN/DailyMail dataset. |
| |
| **Model**: [arinbalyan/summarization-lora](https://huggingface.co/arinbalyan/summarization-lora) |
| """ |
| ) |
|
|
| with gr.Row(): |
| with gr.Column(scale=3): |
| article_input = gr.Textbox( |
| label="Article", |
| placeholder="Paste an article here...", |
| lines=10, |
| ) |
| with gr.Row(): |
| temperature = gr.Slider( |
| minimum=0.1, maximum=1.0, value=0.3, step=0.05, label="Temperature" |
| ) |
| max_tokens = gr.Slider( |
| minimum=50, |
| maximum=250, |
| value=120, |
| step=10, |
| label="Max Summary Tokens", |
| ) |
| summarize_btn = gr.Button("Summarize", variant="primary", size="lg") |
|
|
| with gr.Column(scale=2): |
| summary_output = gr.Textbox( |
| label="Generated Summary", |
| lines=10, |
| interactive=False, |
| ) |
|
|
| with gr.Row(): |
| gr.Markdown("### Try a Sample Article") |
| with gr.Row(): |
| for label, text in SAMPLE_ARTICLES: |
| gr.Button(label, size="sm").click( |
| fn=lambda t=text: t, outputs=article_input |
| ) |
|
|
| summarize_btn.click( |
| fn=summarize, |
| inputs=[article_input, temperature, max_tokens], |
| outputs=summary_output, |
| ) |
|
|
| gr.Markdown( |
| """ |
| --- |
| **Note**: First inference may be slow as the model loads. Subsequent generations |
| are faster. Built with SmolLM2-1.7B fine-tuned via LoRA (r=8, alpha=16) on |
| CNN/DailyMail using a Kaggle P100 GPU. |
| """ |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
