Update app.py

app.py

@@ -60,9 +60,32 @@ def str_to_json(str_obj):
 
 
 
+import subprocess
+
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+
+from accelerate import init_empty_weights, infer_auto_device_map, load_checkpoint_and_dispatch
+from accelerate import Accelerator
+
+
+subprocess.run(
+    "pip install psutil",
+   
+    shell=True,
+)
+
+
+
 # Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("HumanLLMs/Human-Like-Mistral-Nemo-Instruct-2407")
-model = AutoModelForCausalLM.from_pretrained("HumanLLMs/Human-Like-Mistral-Nemo-Instruct-2407", device_map="auto")  # to("cuda:0") 
+model = AutoModelForCausalLM.from_pretrained("HumanLLMs/Human-Like-Mistral-Nemo-Instruct-2407", device_map="auto",
+                                             low_cpu_mem_usage=True,
+    torch_dtype=torch.bfloat16,
+    # quantization_config=quantization_config,
+    attn_implementation="flash_attention_2",
+                                            
+                                            )  # to("cuda:0") 
 terminators = [
     tokenizer.eos_token_id,
     tokenizer.convert_tokens_to_ids("<|eot_id|>")