handler

handler.py

@@ -1,5 +1,6 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftConfig, PeftModel
+from peft import PeftConfig
+from peft import PeftModel
 import torch.cuda
 from typing import Any, Dict
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -7,7 +8,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 class EndpointHandler():
     def __init__(self, path=""):
         config = PeftConfig.from_pretrained(path)
-        model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, load_in_8bit=True, device_map='auto')
+        model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path)
         self.tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
         # Load the Lora model
         self.model = PeftModel.from_pretrained(model, path)