Minor fix

guestrin-lab · Sep 29, 2024 · ecede3f · ecede3f
1 parent d06f68c
commit ecede3f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lotus/models/openai_model.py b/lotus/models/openai_model.py
@@ -52,6 +52,7 @@ def __init__(
         self.use_chat = provider in ["openai", "dbrx", "ollama"]
         self.max_batch_size = max_batch_size
         self.max_ctx_len = max_ctx_len
+        self.hf_name = hf_name if hf_name is not None else model
 
         self.kwargs = {
             "model": model,
@@ -68,7 +69,7 @@ def __init__(
         if self.provider == "openai":
             self.tokenizer = tiktoken.encoding_for_model(model)
         else:
-            self.tokenizer = AutoTokenizer.from_pretrained(hf_name)
+            self.tokenizer = AutoTokenizer.from_pretrained(self.hf_name)
 
     def handle_chat_request(self, messages: List, **kwargs: Dict[str, Any]) -> Union[List, Tuple[List, List]]:
         """Handle single chat request to OpenAI server.