Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 from huggingface_hub import login
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import gradio as gr
 
 # Obtén el token desde la variable de entorno
@@ -12,10 +12,17 @@ if hf_token:
 else:
     raise ValueError("Hugging Face token no encontrado. Asegúrate de que la variable de entorno HF_TOKEN esté configurada.")
 
+# Configuración para cargar el modelo en 4 bits utilizando bitsandbytes
+bnb_config = BitsAndBytesConfig(load_in_4bit=True)
+
 # Cargar el modelo y tokenizador
 model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    quantization_config=bnb_config  # Aplicar cuantización en 4 bits
+)
 
 # Definir la función de inferencia del chatbot
 def chat_fn(multimodal_message):
@@ -70,3 +77,4 @@ with gr.Blocks() as demo:
 if __name__ == "__main__":
     demo.launch()
 
+