fix: use local tokenizer as an option

Author: ChenZiHong-Gavin

graphgen/generate.py

@@ -72,15 +72,18 @@ def main():
         ),
     )
 
+    tokenizer_instance = Tokenizer(model_name=config["tokenizer"])
     synthesizer_llm_client = OpenAIModel(
         model_name=os.getenv("SYNTHESIZER_MODEL"),
         api_key=os.getenv("SYNTHESIZER_API_KEY"),
         base_url=os.getenv("SYNTHESIZER_BASE_URL"),
+        tokenizer_instance=tokenizer_instance,
     )
     trainee_llm_client = OpenAIModel(
         model_name=os.getenv("TRAINEE_MODEL"),
         api_key=os.getenv("TRAINEE_API_KEY"),
         base_url=os.getenv("TRAINEE_BASE_URL"),
+        tokenizer_instance=tokenizer_instance,
     )
 
     graph_gen = GraphGen(
@@ -89,7 +92,7 @@ def main():
         synthesizer_llm_client=synthesizer_llm_client,
         trainee_llm_client=trainee_llm_client,
         search_config=config["search"],
-        tokenizer_instance=Tokenizer(model_name=config["tokenizer"]),
+        tokenizer_instance=tokenizer_instance,
     )
 
     graph_gen.insert(data, config["input_data_type"])

graphgen/models/llm/openai_model.py

@@ -55,6 +55,8 @@ class OpenAIModel(TopkTokenModel):
     rpm: RPM = field(default_factory=lambda: RPM(rpm=1000))
     tpm: TPM = field(default_factory=lambda: TPM(tpm=50000))
 
+    tokenizer_instance: Tokenizer = field(default_factory=Tokenizer)
+
     def __post_init__(self):
         assert self.api_key is not None, "Please provide api key to access openai api."
         self.client = AsyncOpenAI(
@@ -125,8 +127,9 @@ async def generate_answer(
 
         prompt_tokens = 0
         for message in kwargs["messages"]:
-            # TODO: need to use local tokenizer to avoid network call
-            prompt_tokens += len(Tokenizer().encode_string(message["content"]))
+            prompt_tokens += len(
+                self.tokenizer_instance.encode_string(message["content"])
+            )
         estimated_tokens = prompt_tokens + kwargs["max_tokens"]
 
         if self.request_limit: