Wolof-Bridge / app.py
dofbi's picture
update
9a04ccb
raw
history blame
871 Bytes
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()