Update app.py

app.py

@@ -106,16 +106,21 @@ def get_llm_response(prompt: str, model_name: str = "meta-llama/Llama-2-7b-chat-
         max_new_tokens = 1024 if is_detailed_query(prompt) else 256
 
     torch.cuda.empty_cache()
+
+    if not os.path.exists("offload"):
+        os.makedirs("offload")
     
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, token=HF_TOKEN)
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         device_map="auto",
+        offload_folder="offload",  # Specify the folder where weights will be offloaded
         use_safetensors=False,
         trust_remote_code=True,
         torch_dtype=torch.float16,
         token=HF_TOKEN
     )
+
     
     text_gen = pipeline("text-generation", model=model, tokenizer=tokenizer)
     outputs = text_gen(prompt, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.7)