import os
from huggingface_hub import InferenceClient
import time

# Cargar el token desde la variable de entorno
token = os.environ.get("HUGGINGFACE_HUB_TOKEN")
model_name = "google/gemma-2-27b-it"

try:
    generator = pipeline(
        "text-generation", 
        model=model_name, 
        device=0 if torch.cuda.is_available() else -1, 
        use_auth_token=token
    )
    print("Modelo cargado correctamente.")
except Exception as e:
    print(f"Error al cargar el modelo: {e}")
    exit(1)

# Función para procesar la entrada y generar la respuesta
def generate_response(text):
    output = generator(text, max_length=512, num_return_sequences=1)
    response = {
        "choices": [
            {
                "text": output[0]['generated_text'],
                "index": 0,
                "logprobs": None,
                "finish_reason": "stop"
            }
        ],
        "id": "req-12345",  # Reemplazar con un ID único
        "model": model_name,
        "created": int(time.time())
    }
    return response

# Configuración de la interfaz Gradio
import gradio as gr

iface = gr.Interface(
    fn=generate_response,
    inputs="text",
    outputs="json",
    title="API compatible con OpenAI",
    description="Introduce texto para obtener una respuesta del modelo Meta-Llama-3-8B-Instruct."
)

# Ejecutar la interfaz
try:
    iface.launch()
except Exception as e:
    print(f"Error al iniciar la interfaz: {e}")