Spaces:
Running
on
Zero
Running
on
Zero
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() | |