Update app.py

app.py

@@ -1,41 +1,92 @@
 import streamlit as st
-import requests
-import os
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    pipeline,
+    BitsAndBytesConfig
+)
+import torch
 
-# Obtener el token de los secretos
-API_URL = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-Guard-2-8B"
-headers = {"Authorization": f"Bearer {os.getenv('YOUR_HUGGING_FACE_TOKEN')}"}
-
-def query(payload):
-    response = requests.post(API_URL, headers=headers, json=payload)
+# 1. Configuración del Modelo
+@st.cache_resource
+def load_model():
     try:
-        response.raise_for_status()
-        return response.json()
-    except requests.exceptions.HTTPError as err:
-        st.error(f"HTTP error occurred: {err}")
-    except Exception as err:
-        st.error(f"Other error occurred: {err}")
-    return None
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16
+        )
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            "microsoft/Phi-3-mini-4k-instruct",
+            device_map="auto",
+            quantization_config=quantization_config,
+            trust_remote_code=True
+        )
+        
+        tokenizer = AutoTokenizer.from_pretrained(
+            "microsoft/Phi-3-mini-4k-instruct"
+        )
+        
+        return model, tokenizer
+    
+    except Exception as e:
+        st.error(f"Error cargando el modelo: {str(e)}")
+        return None, None
+
+# 2. Interfaz de Streamlit
+st.title("🤖 Chatbot Optimizado para M1")
+st.markdown("Usando Microsoft Phi-3-mini - [Hugging Face](https://huggingface.co/microsoft/Phi-3-mini-4k-instruct)")
+
+# 3. Inicialización de Sesión
+if "messages" not in st.session_state:
+    st.session_state.messages = [
+        {"role": "assistant", "content": "¡Hola! Soy tu asistente AI. ¿En qué puedo ayudarte?"}
+    ]
 
-st.title("LLaMA Chatbot")
-st.subheader("Ask anything to the LLaMA model!")
+# 4. Carga del Modelo
+model, tokenizer = load_model()
 
-user_input = st.text_input("You: ")
-if user_input:
-    output = query({"inputs": user_input})
+# 5. Función de Generación
+def generate_response(prompt):
+    try:
+        messages = [
+            {"role": "user", "content": prompt}
+        ]
+        
+        inputs = tokenizer.apply_chat_template(
+            messages,
+            return_tensors="pt"
+        ).to(model.device)
+        
+        outputs = model.generate(
+            inputs,
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        
+        return tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
     
-    # Imprimir la respuesta completa para depuración
-    st.write("API response:", output)
+    except Exception as e:
+        return f"Error generando respuesta: {str(e)}"
+
+# 6. Interacción del Usuario
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+if prompt := st.chat_input("Escribe tu mensaje..."):
+    # Mostrar input del usuario
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.markdown(prompt)
     
-    if output:
-        # Verificar las posibles claves en la respuesta
-        if "generated_text" in output:
-            response = output["generated_text"]
-        elif isinstance(output, list) and "generated_text" in output[0]:
-            response = output[0]["generated_text"]
-        else:
-            response = "Sorry, I couldn't generate a response."
-    else:
-        response = "Sorry, I couldn't generate a response."
+    # Generar respuesta
+    with st.chat_message("assistant"):
+        with st.spinner("Pensando..."):
+            response = generate_response(prompt)
+            st.markdown(response)
     
-    st.write(f"Chatbot: {response}")
+    st.session_state.messages.append({"role": "assistant", "content": response})