Update app.py

app.py

@@ -1,24 +1,22 @@
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+import requests
+import os
 
-# Cargar el modelo y el tokenizador
-@st.cache(allow_output_mutation=True)
-def load_model():
-    model_name = "meta-llama/Meta-Llama-Guard-2-8B"  # Reemplaza con el nombre del modelo que vas a usar
-    use_auth_token = 'YOUR_HUGGING_FACE_TOKEN'
-    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=use_auth_token)
-    model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=use_auth_token)
-    return tokenizer, model
+# 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')}"}
 
-tokenizer, model = load_model()
+def query(payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.json()
 
 st.title("LLaMA Chatbot")
 st.subheader("Ask anything to the LLaMA model!")
 
 user_input = st.text_input("You: ")
 if user_input:
-    inputs = tokenizer(user_input, return_tensors="pt")
-    outputs = model.generate(inputs.input_ids, max_length=150, num_return_sequences=1)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    output = query({"inputs": user_input})
+    response = output.get("generated_text", "Sorry, I couldn't generate a response.")
     st.write(f"Chatbot: {response}")