from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr

# Charger le modèle
model_name = "soynade-research/Oolel-v0.1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")

# Fonction pour générer une réponse
def generate_response(user_input, max_new_tokens=150, temperature=0.7):
    inputs = tokenizer(user_input, return_tensors="pt").to("cuda")
    outputs = model.generate(inputs.input_ids, max_new_tokens=max_new_tokens, temperature=temperature)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Interface Gradio
iface = gr.Interface(
    fn=generate_response,
    inputs=[gr.Textbox(label="Message utilisateur"), gr.Slider(50, 500, value=150, label="Nombre max de tokens")],
    outputs="text",
    title="Oolel Chatbot"
)

iface.launch()