test

app.py

@@ -26,9 +26,6 @@ from transformers import LlamaTokenizer
 from transformers import PreTrainedModel
 from transformers import PreTrainedTokenizerBase
 
-transformers.AutoTokenizer.from_pretrained = LlamaTokenizer.from_pretrained
-transformers.AutoModelForCausalLM.from_pretrained = LlamaForCausalLM.from_pretrained
-
 
 def load_lora_model(
     model_path: str,
@@ -41,13 +38,13 @@ def load_lora_model(
 ) -> Tuple[Union[PreTrainedModel, PeftModel], PreTrainedTokenizerBase]:
     model: Union[PreTrainedModel, PeftModel]
     tokenizer: PreTrainedTokenizerBase
-    model, tokenizer = load_model(
-        model_path=model_path,
-        device=device,
-        num_gpus=num_gpus,
-        max_gpu_memory=max_gpu_memory,
-        load_8bit=load_8bit,
-        debug=debug,
+    tokenizer = LlamaTokenizer.from_pretrained(model_path)
+    model = LlamaForCausalLM.from_pretrained(
+        model_path,
+        load_in_8bit=load_8bit,
+        device_map=device,
+        max_memory=max_gpu_memory,
+        torch_dtype=torch.float16,
     )
     if lora_weight is not None:
         # model = PeftModelForCausalLM.from_pretrained(model, model_path, **kwargs)