Spaces:
Paused
Paused
File size: 2,226 Bytes
f29521b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import gradio as gr
from transformers import pipeline, login
import torch
from dotenv import load_dotenv
import os
import spaces # Importar la librer铆a de spaces
# Cargar variables de entorno desde el archivo .env
load_dotenv()
# Obtener el token de Hugging Face desde las variables de entorno
huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
# Verificar si el token fue cargado correctamente
if not huggingface_token:
raise ValueError("Token de Hugging Face no encontrado en las variables de entorno")
# Iniciar sesi贸n con tu token de Hugging Face
login(token=huggingface_token) # Inicia sesi贸n con el token
# Funci贸n para cargar el modelo en GPU (si est谩 disponible) o en CPU si hay un error
def load_model():
model_id = "facebook/MobileLLM-R1.5-950M"
try:
# Intentar cargar el modelo en GPU
print("Intentando cargar el modelo en la GPU...")
return pipeline(
"text-generation",
model=model_id,
torch_dtype="auto",
device_map="auto", # Esto autom谩ticamente usa la GPU si est谩 disponible
)
except Exception as e:
# Si hay un error (por ejemplo, no hay GPU disponible), usar CPU
print("Error en GPU, cambiando a CPU. Error:", e)
return pipeline(
"text-generation",
model=model_id,
torch_dtype="auto",
device=0, # Usar CPU en caso de error
)
# Cargar el modelo con el fallback a CPU
pipe = load_model()
# Decorador para garantizar que el c贸digo se ejecute en GPU si est谩 disponible
@spaces.GPU
def chat_with_model(system_message, user_message):
messages = [
{"role": "system", "content": system_message},
{"role": "user", "content": user_message},
]
# Generar la respuesta usando el modelo
outputs = pipe(messages, max_new_tokens=8192)
return outputs[0]["generated_text"][-1]
# Crear la interfaz de Gradio
gr.Interface(
fn=chat_with_model,
inputs=[
gr.Textbox(label="System Message", placeholder="Escribe el mensaje del sistema..."),
gr.Textbox(label="User Message", placeholder="Escribe tu mensaje aqu铆..."),
],
outputs="text",
live=True,
).launch()
|