Update app.py

app.py

@@ -9,6 +9,13 @@ import re
 from transformers import pipeline
 from uuid import uuid4
 
+import os
+import tempfile
+
+offload_dir='/content/offload'
+os.makedirs(offload_dir) if not os.path.exists(offload_dir) else None
+
+
 #based on https://huggingface.co/spaces/Wauplin/space_to_dataset_saver/blob/main/app_json.py
 #data is saved at https://huggingface.co/datasets/MR17u/tweeteval-irony-mcc/tree/main
 
@@ -33,7 +40,7 @@ Classify if the following tweet is ironic or not
 #     path_in_repo="data",
 # )
 
-classifier = pipeline("text-generation", model="meta-llama/Llama-2-7b-hf", low_cpu_mem_usage=True)
+classifier = pipeline("text-generation", model="meta-llama/Llama-2-7b-hf", low_cpu_mem_usage=True, device_map="auto", offload_folder=offload_dir)
 classifier.load_lora_weights("PierreEpron/llama7b-irony", weight_name="adapter_model.safetensors")
 
 def clean_brackets(text):