Update handler.py

handler.py

@@ -1,4 +1,3 @@
-from optimum.onnxruntime import ORTModelForCausalLM
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import re
 import time
@@ -23,9 +22,14 @@ Alice Gate: *Alice strides into the room with a smile, her eyes lighting up when
 
 class SweetCommander():
 
-    def __init__(self, path="BlueDice/Katakuri-350m-onnx") -> None:
+    def __init__(self, path="BlueDice/Katakuri-350m") -> None:
         self.tokenizer = AutoTokenizer.from_pretrained(path)
-        self.model = ORTModelForCausalLM.from_pretrained(path)#provider = "CUDAExecutionProvider"
+        self.model = AutoModelForCausalLM.from_pretrained(
+            path,
+            low_cpu_mem_usage = True, 
+            trust_remote_code = False,
+            torch_dtype = torch.float32,
+        )
         self.star_line = "***********************************************************"
 
     def __call__(self, user_name, user_input):