Update app.py

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 import json
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import os
 
 # Recuperar el token de los secrets de Hugging Face Spaces
@@ -9,27 +9,32 @@ hf_token = os.getenv('hf_token')
 
 # Configuración del modelo LLaMA
 model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
 
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     torch_dtype=torch.bfloat16,
     device_map="auto",
+    quantization_config=bnb_config,
     token=hf_token
 )
 
 SYS_PROMPT = """You are an AI assistant tasked with analyzing the sentiment of a given text at the aspect level, specifically related to public security issues in Ecuador, such as crime, violence, law enforcement, or corruption. For each aspect mentioned in the text, classify the sentiment as 'Positive', 'Neutral', or 'Negative' and provide a clear justification for each aspect. Return the response **only** in JSON format and in Spanish, without additional information."""
 
 def generate(formatted_prompt):
-    # Crear un único input para tokenizar
-    input_text = SYS_PROMPT + formatted_prompt
-
-    # Tokenizar el texto
-    input_ids = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True).input_ids.to(model.device)
+    messages = [{"role": "system", "content": SYS_PROMPT}, {"role": "user", "content": formatted_prompt}]
 
-    # Verificar si input_ids no es None
-    if input_ids is None or input_ids.size(0) == 0:
-        return {"error": "No se pudo generar input_ids"}
+    input_ids = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        return_tensors="pt"
+    ).to(model.device)
 
     attention_mask = input_ids.ne(tokenizer.pad_token_id)