Update app.py

app.py

@@ -12,6 +12,8 @@ from transformers import (
     TextStreamer
 )
 
+from accelerate import disk_offload
+
 
 model = AutoModelForCausalLM.from_pretrained(
     'DuongTrongChi/Rikka-1.8B-v2.2',
@@ -19,6 +21,9 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
+disk_offload(model=model, offload_dir="offload")
+
+
 tokenizer = AutoTokenizer.from_pretrained('DuongTrongChi/Rikka-1.8B-v2.2')
 eos_token = tokenizer("<|im_end|>",add_special_tokens=False)["input_ids"][0]
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)