app_1.py

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import gradio as gr
import os

# Asegúrate de que tu token de Hugging Face está cargado como una variable de entorno
hf_token = os.environ.get("token")

if hf_token is not None:
    from huggingface_hub import HfFolder
    HfFolder.save_token(hf_token)
else:
    print("No se encontró el token de Hugging Face. Asegúrate de que la variable de entorno HF_TOKEN esté configurada.")

# Cargar el tokenizador
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")

# Añadir el token especial [PAD]
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
print("Token [PAD] añadido al tokenizador.")

# Asegurarse de que el modelo base usa el tamaño de vocabulario actualizado
model_base = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it")
model_base.resize_token_embeddings(len(tokenizer))

# Cargar el modelo con el adaptador
# Nota: Asegúrate de que el enfoque utilizado para cargar el adaptador es compatible
# con la manera en que guardaste el adaptador y su configuración en Hugging Face Hub.
model_with_adapter = PeftModel.from_pretrained(model_base, "Juliofc/chaterapia_model")



def generate_text(input_text):
    # Codificar el texto de entrada
    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)

    # Generar respuesta del modelo
    output_ids = model_with_adapter.generate(input_ids, max_length=50, num_return_sequences=1)

    # Decodificar y retornar la respuesta
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

# Crear la interfaz de Gradio
gradio_app = gr.Interface(fn=generate_text, 
                     inputs=gr.Textbox(lines=2, placeholder="Escribe aquí..."), 
                     outputs="text", 
                     title="Generador de Texto con GEMMA 2B",
                     description="Modelo GEMMA 2B con adaptador para generación de texto.")

if __name__ == "__main__":
    gradio_app.launch()