Update app.py

app.py

@@ -1,25 +1,21 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-import os
-from huggingface_hub import login
 
 # Inicializar el modelo generador de texto
 @st.cache_resource  # Cachear el modelo para no recargarlo repetidamente
 def load_model():
-    read_key = os.getenv("HUGGINGFACE_TOKEN")
-    login(token=read_key)
-    
     model_id = "meta-llama/Llama-2-7b"  # Cambia al modelo correcto disponible
-    return pipeline(
-        "text-generation",
-        model=model_id,
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
         torch_dtype=torch.bfloat16,
         device_map="auto"
     )
+    return tokenizer, model
 
-# Cargar el modelo
-pipe = load_model()
+# Cargar el modelo y el tokenizador
+tokenizer, model = load_model()
 
 # Crear la interfaz de usuario
 st.title("Generador de Texto Llama-2")
@@ -31,13 +27,22 @@ user_input = st.text_input("Escribe tu prompt aquí:", "")
 # Botón para generar texto
 if st.button("Generar texto"):
     if user_input:
-        # Generar el texto utilizando el modelo
+        # Tokenizar el input y generar el texto
+        input_ids = tokenizer.encode(user_input, return_tensors='pt').to(model.device)
         with st.spinner("Generando..."):
-            result = pipe(user_input, do_sample=True, max_length=100)
-            generated_text = result[0]['generated_text']
+            output = model.generate(
+                input_ids,
+                do_sample=True,
+                max_length=100,
+                num_return_sequences=1,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        # Decodificar el texto generado
+        generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
         
         # Mostrar el texto generado
         st.subheader("Texto generado:")
         st.write(generated_text)
     else:
-        st.error("Por favor, introduce un prompt antes de generar texto.")
+        st.error("Por favor, introduce un prompt antes de generar texto.")