import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Cargar el tokenizer del modelo base
tokenizer = AutoTokenizer.from_pretrained("unsloth/meta-llama-3.1-8b-bnb-4bit")

# Cargar el modelo base sin usar bitsandbytes ni 4-bit
model = AutoModelForCausalLM.from_pretrained("unsloth/meta-llama-3.1-8b-bnb-4bit", device_map="cpu")

# Cargar el adaptador LoRA
model.load_adapter("JuanCabs/lapepav1", source="hf")

# Función para generar texto con el modelo
def generar_respuesta(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=100)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Crear la interfaz de Gradio
interface = gr.Interface(fn=generar_respuesta, inputs="text", outputs="text", 
                         title="Generación de Texto con Lora y Llama 3.1",
                         description="Introduce un texto y genera una respuesta usando un modelo con LoRA.")

# Lanzar la aplicación en Hugging Face Spaces
if __name__ == "__main__":
    interface.launch()