Update code/inference.py

code/inference.py

@@ -1,4 +1,4 @@
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import re
 
@@ -21,11 +21,15 @@ Alice Gate: *Alice strides into the room with a smile, her eyes lighting up when
 Alice Gate:"""
 
 def model_fn(model_dir):
-  # Load model from HuggingFace Hub
-  tokenizer = AutoTokenizer.from_pretrained(model_dir)
-  model = torch.load(f"{model_dir}/torch_model.pt")
-  return model, tokenizer
-
+    # Load model from HuggingFace Hub
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_dir,
+        low_cpu_mem_usage = True, 
+        trust_remote_code = False,
+        torch_dtype = torch.float16,
+    ).to('cuda')
+    return model, tokenizer
 
 def create_new_response(result, user_name):
     result = result.rsplit("Alice Gate:", 1)[1].split(f"{user_name}:",1)[0].strip()