owen/worker-vllm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix tokenizer's trust_remote_code parameter

Browse Source
This commit is contained in:
alpayariyak
2024-02-08 23:53:13 +00:00
parent 2b5b8dfb61
commit 4f5e0d37c4
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine.py
+3 -3
View File
@@ -13,8 +13,8 @@ from dotenv import load_dotenv
class Tokenizer: class Tokenizer:
def __init__(self, tokenizer_name_or_path, tokenizer_revision): def __init__(self, tokenizer_name_or_path, tokenizer_revision, trust_remote_code):
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, revision=tokenizer_revision) self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, revision=tokenizer_revision, trust_remote_code=trust_remote_code)
self.custom_chat_template = os.getenv("CUSTOM_CHAT_TEMPLATE") self.custom_chat_template = os.getenv("CUSTOM_CHAT_TEMPLATE")
self.has_chat_template = bool(self.tokenizer.chat_template) or bool(self.custom_chat_template) self.has_chat_template = bool(self.tokenizer.chat_template) or bool(self.custom_chat_template)
if self.custom_chat_template and isinstance(self.custom_chat_template, str): if self.custom_chat_template and isinstance(self.custom_chat_template, str):
@@ -41,7 +41,7 @@ class vLLMEngine:
load_dotenv() # For local development load_dotenv() # For local development
self.config = self._initialize_config() self.config = self._initialize_config()
logging.info("vLLM config: %s", self.config) logging.info("vLLM config: %s", self.config)
self.tokenizer = Tokenizer(self.config["tokenizer"], self.config["tokenizer_revision"]) self.tokenizer = Tokenizer(self.config["tokenizer"], self.config["tokenizer_revision"], self.config["trust_remote_code"])
self.llm = self._initialize_llm() if engine is None else engine self.llm = self._initialize_llm() if engine is None else engine
self.openai_engine = self._initialize_openai() self.openai_engine = self._initialize_openai()
self.max_concurrency = int(os.getenv("MAX_CONCURRENCY", DEFAULT_MAX_CONCURRENCY)) self.max_concurrency = int(os.getenv("MAX_CONCURRENCY", DEFAULT_MAX_CONCURRENCY))