Update app.py

app.py

@@ -1,4 +1,5 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch, dispatch_model, infer_auto_device_map
 import streamlit as st
 from huggingface_hub import login
 import pandas as pd
@@ -10,13 +11,23 @@ login(huggingface_token)
 # Cargar el tokenizador y el modelo
 model_id = "meta-llama/Llama-3.2-1B"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(model_id) #, device_map="auto")
 tokenizer.pad_token = tokenizer.eos_token
 
 MAX_INPUT_TOKEN_LENGTH = 10000
 
+# Cargar el modelo con disk_offload
+with init_empty_weights():
+    model = AutoModelForCausalLM.from_config(model_id)
+
+device_map = infer_auto_device_map(model, max_memory={"disk": "2GiB"}, no_split_module_classes=["LlamaDecoderLayer"])
+model = load_checkpoint_and_dispatch(model, model_id, device_map=device_map, offload_folder="offload_dir")
+
+MAX_INPUT_TOKEN_LENGTH = 10000
+
+
 def generate_response(input_text, temperature=0.7, max_new_tokens=20):
-    input_ids = tokenizer.encode(input_text, return_tensors='pt').to(model.device)
+    input_ids = tokenizer.encode(input_text, return_tensors='pt').to("cpu")  # Usar 'cpu' para mantener la compatibilidad
 
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
@@ -32,7 +43,7 @@ def generate_response(input_text, temperature=0.7, max_new_tokens=20):
         top_p=0.9,
         temperature=temperature,
         num_return_sequences=3,
-        eos_token_id=tokenizer.eos_token_id  # Cambiado a un entero
+        eos_token_id=tokenizer.eos_token_id
     )
 
     try: