inference

config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "imone/LLaMA2_13B_with_EOT_token",
+  "_name_or_path": "swapnice/swapnice-openorcaxopenchat-preview2-13b",
   "architectures": [
     "LlamaForCausalLM"
   ],

handler.py

@@ -3,17 +3,16 @@ from typing import Dict, List, Any
 from transformers import LlamaForCausalLM, LlamaTokenizer, pipeline
 
 # get dtype
-dtype = torch.float16
+dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
 
 class EndpointHandler:
     def __init__(self, path=""):
         # load the model
-        tokenizer = LlamaTokenizer.from_pretrained(".", local_files_only=True)
+        tokenizer = LlamaTokenizer.from_pretrained("/repository/tokenizer", local_files_only=True)
         model = LlamaForCausalLM.from_pretrained(
-            ".",
+            "/repository/pytorch_model",
             device_map="auto",
             torch_dtype=dtype,
-            offload_folder="offload"
         )
         # create inference pipeline
         self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)