Update app.py

app.py

@@ -1,16 +1,21 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 import requests
 import os
 
 # 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')}"}
+headers = {"Authorization": f"Bearer {os.getenv('HF_API_TOKEN')}"}
 
 def query(payload):
     response = requests.post(API_URL, headers=headers, json=payload)
-    return response.json()
+    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
 
 st.title("LLaMA Chatbot")
 st.subheader("Ask anything to the LLaMA model!")
@@ -18,5 +23,19 @@ st.subheader("Ask anything to the LLaMA model!")
 user_input = st.text_input("You: ")
 if user_input:
     output = query({"inputs": user_input})
-    response = output.get("generated_text", "Sorry, I couldn't generate a response.")
+    
+    # Imprimir la respuesta completa para depuración
+    st.write("API response:", output)
+    
+    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."
+    
     st.write(f"Chatbot: {response}")