import os, torch, gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Use Space env var MODEL_ID if set; otherwise fall back to default below
MODEL_ID = os.getenv("MODEL_ID", "Shahzeb99/Article_Summarizer")
HF_TOKEN = os.getenv("HF_TOKEN")  # add as a secret if model is private

device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True, token=HF_TOKEN)
model = AutoModelForSeq2SeqLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16 if torch.cuda.is_available() else None,
    device_map="auto" if torch.cuda.is_available() else None,
    token=HF_TOKEN
).to(device)

def summarize_fn(text, max_new_tokens, min_new_tokens, num_beams, length_penalty):
    if not text or not text.strip():
        return ""
    enc = tokenizer("highlights: " + text.strip(), return_tensors="pt", truncation=True, max_length=512)
    enc = {k: v.to(model.device) for k, v in enc.items()}
    with torch.no_grad():
        out = model.generate(
            **enc,
            max_new_tokens=int(max_new_tokens),
            min_new_tokens=int(min_new_tokens),
            num_beams=int(num_beams),
            length_penalty=float(length_penalty),
            early_stopping=True,
            no_repeat_ngram_size=3
        )
    return tokenizer.decode(out[0], skip_special_tokens=True)

with gr.Blocks(title="Article → Highlights") as demo:
    tip = "" if "Shahzeb99/Article_Summarizer" not in MODEL_ID else "⚠️ Set MODEL_ID in Space settings or edit app.py."
    gr.Markdown("## Article → Highlights\n" + tip)
    inp = gr.Textbox(lines=12, label="Article")
    max_new = gr.Slider(32, 512, value=150, step=8, label="Max new tokens")
    min_new = gr.Slider(8, 200, value=40, step=4, label="Min new tokens")
    beams = gr.Slider(1, 8, value=4, step=1, label="Beam size")
    lp = gr.Slider(0.2, 3.0, value=2.0, step=0.1, label="Length penalty")
    btn = gr.Button("Generate")
    out = gr.Textbox(lines=8, label="Highlights")
    btn.click(summarize_fn, [inp, max_new, min_new, beams, lp], out)

# works across recent Gradio versions
demo.queue()   # or just remove this line entirely
if __name__ == "__main__":
    demo.launch(share=True)