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 for quick testing
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()