import streamlit as st
import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoModelForSeq2SeqLM
from peft import PeftModel
from newspaper import Article

st.title("LORYX: LoRA vs Baseline Summarizer")

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

# -------- LOAD MODELS --------
@st.cache_resource
def load_model():
    base_model = T5ForConditionalGeneration.from_pretrained("t5-base")
    tokenizer = T5Tokenizer.from_pretrained("t5-base")

    lora_model = PeftModel.from_pretrained(base_model, "t5_lora_news_adapter")
    lora_model.to(device).eval()

    baseline_model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")
    baseline_model.to(device).eval()

    return lora_model, baseline_model, tokenizer

lora_model, baseline_model, tokenizer = load_model()

# -------- EXTRACT TEXT FROM URL --------
def extract_text(url):
    try:
        article = Article(url)
        article.download()
        article.parse()
        return article.text
    except:
        return None

# -------- SUMMARIZATION FUNCTION --------
def generate_summary(model, text):
    inputs = tokenizer(
        "summarize: " + text,
        return_tensors="pt",
        max_length=512,
        truncation=True
    ).to(device)

    output = model.generate(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_length=200,
        min_length=100,
        num_beams=8,
        length_penalty=2.0,
        no_repeat_ngram_size=3,
        early_stopping=True
    )

    return tokenizer.decode(output[0], skip_special_tokens=True)

# -------- INPUT SECTION --------
option = st.radio("Choose Input Type", ["Text", "URL"])

if option == "Text":
    user_input = st.text_area("Enter Article Text")

else:
    user_input = st.text_input("Paste Article URL")

# -------- BUTTON --------
if st.button("Generate Summary"):

    if not user_input:
        st.warning("Please provide input")
    else:
        # Handle URL
        if option == "URL":
            text = extract_text(user_input)

            if not text:
                st.error("❌ Failed to extract article")
                st.stop()
        else:
            text = user_input

        text = text[:2000]  # important truncation

        # Generate both summaries
        lora_summary = generate_summary(lora_model, text)
        baseline_summary = generate_summary(baseline_model, text)

        # Output
        st.subheader("Fine-Tuned (LoRA) Summary")
        st.write(lora_summary)

        st.subheader("Baseline Summary")
        st.write(baseline_summary)