from transformers import pipeline
import gradio as gr
import joblib

# Configuration de base
# model_name = "facebook/bart-base"
# output_dir = "./bart_blagueur_model"
# # Tokenizer et modèle
# tokenizer = BartTokenizer.from_pretrained(model_name)
# model = BartForConditionalGeneration.from_pretrained(model_name)
model__ = joblib.load("model.joblib")
tokenizer_ = joblib.load("tokenizer.joblib")

def blagueur(prompt):
    inputs = tokenizer_(prompt, return_tensors="pt", truncation=True, padding="max_length", max_length=128)
    outputs = model__.generate(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_length=64,
        num_beams=5,
        do_sample=True,
        temperature=0.9)
    return tokenizer_.decode(outputs[0], skip_special_tokens=True)
# Configuration de l'interface Gradio
demo = gr.Interface(
    fn=blagueur,
    inputs=gr.Textbox(lines=2, placeholder="Demandez-moi une blague..."),
    outputs="text",
    title="Chatbot Comique",
    description="Un chatbot qui raconte des blagues. Demandez-lui une blague et il vous fera rire!",
    examples=[
        ["Raconte-moi une blague"],
        ["Dis-moi une blague sur les animaux"],
        ["Blague sur les informaticiens"]
    ],
    theme='shivi/calm_seafoam'
)

demo.launch()