Fix memory management and offload dirs

Author: RETR0-OS

ModelForge/utilities/chat_playground.py

@@ -39,6 +39,7 @@ def __init__(self, model_path: str):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         if self.device == "cpu":
             print("CUDA is not available. Exiting...")
+            input("\nPress Enter to exit...")
             exit(1)
 
         print("Loading model...")
@@ -52,7 +53,15 @@ def __init__(self, model_path: str):
             streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
             tokenizer.pad_token = tokenizer.eos_token
             module = getattr(peft, self.modelforge_config["model_class"])
-            peft_model = module.from_pretrained(model_path, config=config, is_trainable=False)
+            peft_model = module.from_pretrained(
+                model_path,
+                config=config,
+                is_trainable=False,
+                low_cpu_mem_usage=True,
+                device_map="auto",
+                torch_dtype="auto",
+                offload_folder=os.path.join(model_path, "offload")
+            )
             self.generator = pipeline(
                 self.modelforge_config["pipeline_task"],
                 streamer=streamer,
@@ -61,12 +70,15 @@ def __init__(self, model_path: str):
             )
         except AttributeError:
             print(f"Model class {self.modelforge_config['model_class']} not found in peft module.")
+            input("\nPress Enter to exit...")
             exit(1)
         except KeyError:
             print(f"Pipeline task {self.modelforge_config['pipeline_task']} not found in definitions for the pipeline object of transformers.")
+            input("\nPress Enter to exit...")
             exit(1)
         except Exception as e:
             print(traceback.format_exc())
+            input("\nPress Enter to exit...")
             exit(1)
 
     def generate_response(self, prompt: str, context=None, temperature=0.2, top_p=0.92, top_k=50,