Update app.py

app.py

@@ -5,21 +5,17 @@ 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'
+# Set the Hugging Face home directory
+os.environ['HF_HOME'] = '/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',
-    cache_dir='/app/.cache'
+    device_map='auto'
 )
 
 # Load the tokenizer
-tokenizer = AutoTokenizer.from_pretrained(
-    "SHASWATSINGH3101/Qwen2-0.5B-Instruct_lora_merge",
-    cache_dir='/app/.cache'
-)
+tokenizer = AutoTokenizer.from_pretrained("SHASWATSINGH3101/Qwen2-0.5B-Instruct_lora_merge")
 tokenizer.pad_token = tokenizer.eos_token
 
 def gen(model, p, maxlen=100, sample=True):
@@ -49,7 +45,8 @@ iface = gr.Interface(
     inputs=gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
     outputs="text",
     title="Legal Letter Generator",
-    description="Generate a letter informing someone of potential legal action due to a dispute or violation."
+    description="Generate a letter informing someone of potential legal action due to a dispute or violation.",
+    flagging_dir="/app/flagged"  # Set the flagging directory
 )
 
 # Launch the app