import gradio as gr
import spaces
import torch
from transformers import pipeline

# Prueba de disponibilidad de GPU
zero = torch.Tensor([0]).cuda()
print(zero.device)  # Imprime el dispositivo actual, por ejemplo 'cpu' o 'cuda:0'

# Inicializa el pipeline de generación de texto en GPU (device=0)
pipe = pipeline("text-generation", model="aaditya/Llama3-OpenBioLLM-70B", device=0)

@spaces.GPU
def generate_response(prompt):
    print(zero.device)  # Ahora debería imprimir 'cuda:0'
    # Genera texto a partir del prompt (ajusta max_length u otros parámetros según necesidad)
    result = pipe(prompt, max_length=100)
    return result[0]['generated_text']

# Configuración de la interfaz de Gradio
iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(label="Introduce tu prompt", placeholder="Escribe aquí tu pregunta o indicación..."),
    outputs=gr.Textbox(label="Respuesta generada"),
    title="Generador de Texto con Llama3-OpenBioLLM-70B",
    description="Ingresa un prompt y obtén una respuesta generada por el modelo Llama3-OpenBioLLM-70B."
)

iface.launch()