8B-on-GPU-ZERO / app.py
alex-abb's picture
Update app.py
896f309 verified
raw
history blame
790 Bytes
import spaces
import gradio as gr
from transformers import pipeline
import torch
print(f"Is CUDA available: {torch.cuda.is_available()}")
# True
print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
# Tesla T4
# Initialiser le générateur de texte avec GPT-2
generator = pipeline('text-generation', model='gpt2')
# Optionnel : Fixer une graine aléatoire pour la reproductibilit
# Fonction de génération de texte
@spaces.gpu
def generate_text(prompt):
response = generator(prompt, max_length=100, num_return_sequences=1, truncation=True)
return response[0]['generated_text']
# Définir une fonction pour l'interface de chat
def chatbot(message,history):
str(message)
return generate_text(message)
gr.ChatInterface(chatbot).launch()