Update app.py

app.py

@@ -1,17 +1,25 @@
+import os
 import torch
 import pandas as pd
 import numpy as np
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
 
+# Set the cache directory
+os.environ['TRANSFORMERS_CACHE'] = '/app/.cache'
+
 # Load the base model with device_map set to 'auto'
 model = AutoModelForCausalLM.from_pretrained(
     "SHASWATSINGH3101/Qwen2-0.5B-Instruct_lora_merge",
-    device_map='auto'
+    device_map='auto',
+    cache_dir='/app/.cache'
 )
 
 # Load the tokenizer
-tokenizer = AutoTokenizer.from_pretrained("SHASWATSINGH3101/Qwen2-0.5B-Instruct_lora_merge")
+tokenizer = AutoTokenizer.from_pretrained(
+    "SHASWATSINGH3101/Qwen2-0.5B-Instruct_lora_merge",
+    cache_dir='/app/.cache'
+)
 tokenizer.pad_token = tokenizer.eos_token
 
 def gen(model, p, maxlen=100, sample=True):