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

# Load model and tokenizer explicitly
model_name = "facebook/mbart-large-50"
tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang="ne_NP")
model = AutoModelForSeq2SeqLM.from_pretrained(
    model_name,
    device_map="auto",
    low_cpu_mem_usage=True
)

summarizer = pipeline(
    "summarization",
    model=model,
    tokenizer=tokenizer
)

def summarize_text(text):
    try:
        if not text.strip():
            return "Please enter some Nepali text to summarize"
            
        summary = summarizer(
            text,
            max_length=1000,
            min_length=30,
            truncation=True,
            # Directly pass forced_bos_token_id here
            forced_bos_token_id=tokenizer.lang_code_to_id["ne_NP"]
        )[0]['summary_text']
        return summary
    except Exception as e:
        return f"Error during summarization: {str(e)}"

iface = gr.Interface(
    fn=summarize_text,
    inputs=gr.Textbox(lines=5, label="Nepali Text to Summarize"),
    outputs=gr.Textbox(lines=5, label="Summary"),
    title="Nepali Text Summarizer",
    description="Enter Nepali text and get a concise summary using multilingual NLP models.",
    allow_flagging="never"
)

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)