8B-on-GPU-ZERO / app.py
alex-abb's picture
Update app.py
42fc494 verified
raw
history blame
889 Bytes
import gradio as gr
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", device_map="auto")
# Initialiser le générateur de texte avec GPT-2
# Optionnel : Fixer une graine aléatoire pour la reproductibilit
@spaces.GPU(duration=120)
# Fonction de génération de texte
def generate_text(prompt):
inputs = tokenizer(prompt, return_tensors="pt")
response_ids = model.generate(inputs.input_ids)
response_text = tokenizer.decode(response_ids[0], skip_special_tokens=True)
return response_text
# Définir une fonction pour l'interface de chat
def chatbot(message,history):
return generate_text(message)
gr.ChatInterface(chatbot).launch()