import gradio as gr
from transformers import BartTokenizer, BartForConditionalGeneration

# Load your model and tokenizer from the Hugging Face Hub
model_name = "bart_headline_model"  # Replace with your uploaded model name
tokenizer = BartTokenizer.from_pretrained(model_name)
#model = BartForConditionalGeneration.from_pretrained(model_name)
model = BartForConditionalGeneration.from_pretrained(model_name, device_map="cpu")


# Define the function for generating captions
def generate_caption(article):
    inputs = tokenizer(article, return_tensors="pt", max_length=128, truncation=True)
    outputs = model.generate(inputs["input_ids"], max_length=20, num_beams=5, early_stopping=True)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Create the Gradio interface
interface = gr.Interface(
    fn=generate_caption,
    inputs="text",
    outputs="text",
    title="Headline Generator",
    description="Enter an article to generate a headline.",
)

# Launch the app
interface.launch()