fix: change variable name

Author: ChenZiHong-Gavin

generate.py

@@ -48,7 +48,7 @@ def save_config(global_config):
         api_key=os.getenv("SYNTHESIZER_API_KEY"),
         base_url=os.getenv("SYNTHESIZER_BASE_URL")
     )
-    training_llm_client = OpenAIModel(
+    trainee_llm_client = OpenAIModel(
         model_name=os.getenv("TRAINEE_MODEL"),
         api_key=os.getenv("TRAINEE_API_KEY"),
         base_url=os.getenv("TRAINEE_BASE_URL")
@@ -61,7 +61,7 @@ def save_config(global_config):
     graph_gen = GraphGen(
         unique_id=unique_id,
         synthesizer_llm_client=synthesizer_llm_client,
-        training_llm_client=training_llm_client,
+        trainee_llm_client=trainee_llm_client,
         if_web_search=config['web_search'],
         tokenizer_instance=Tokenizer(
             model_name=config['tokenizer']

graphgen/graphgen.py

@@ -29,7 +29,7 @@ class GraphGen:
 
     # llm
     synthesizer_llm_client: OpenAIModel = None
-    training_llm_client: OpenAIModel = None
+    trainee_llm_client: OpenAIModel = None
     tokenizer_instance: Tokenizer = None
 
     # web search
@@ -181,7 +181,7 @@ def judge(self, re_judge=False):
         loop.run_until_complete(self.async_judge(re_judge))
 
     async def async_judge(self, re_judge=False):
-        _update_relations = await judge_statement(self.training_llm_client, self.graph_storage,
+        _update_relations = await judge_statement(self.trainee_llm_client, self.graph_storage,
                                                   self.rephrase_storage, re_judge)
         await _update_relations.index_done_callback()
 

graphgen/operators/judge.py

@@ -7,15 +7,15 @@
 
 
 async def judge_statement( # pylint: disable=too-many-statements
-        training_llm_client: OpenAIModel,
+        trainee_llm_client: OpenAIModel,
         graph_storage: NetworkXStorage,
         rephrase_storage: JsonKVStorage,
         re_judge: bool = False,
         max_concurrent: int = 1000) -> NetworkXStorage:
     """
     Get all edges and nodes and judge them
 
-    :param training_llm_client: judge the statements to get comprehension loss
+    :param trainee_llm_client: judge the statements to get comprehension loss
     :param graph_storage: graph storage instance
     :param rephrase_storage: rephrase storage instance
     :param re_judge: re-judge the relations
@@ -46,7 +46,7 @@ async def _judge_single_relation(
                 judgements = []
                 gts = [gt for _, gt in descriptions]
                 for description, gt in descriptions:
-                    judgement = await training_llm_client.generate_topk_per_token(
+                    judgement = await trainee_llm_client.generate_topk_per_token(
                         STATEMENT_JUDGEMENT_PROMPT['TEMPLATE'].format(statement=description)
                     )
                     judgements.append(judgement[0].top_candidates)
@@ -94,7 +94,7 @@ async def _judge_single_entity(
                 judgements = []
                 gts = [gt for _, gt in descriptions]
                 for description, gt in descriptions:
-                    judgement = await training_llm_client.generate_topk_per_token(
+                    judgement = await trainee_llm_client.generate_topk_per_token(
                         STATEMENT_JUDGEMENT_PROMPT['TEMPLATE'].format(statement=description)
                     )
                     judgements.append(judgement[0].top_candidates)

webui/app.py

@@ -43,7 +43,7 @@ def init_graph_gen(config: dict, env: dict) -> GraphGen:
         api_key=env.get("SYNTHESIZER_API_KEY", "")
     )
 
-    graph_gen.training_llm_client = OpenAIModel(
+    graph_gen.trainee_llm_client = OpenAIModel(
         model_name=env.get("TRAINEE_MODEL", ""),
         base_url=env.get("TRAINEE_BASE_URL", ""),
         api_key=env.get("TRAINEE_API_KEY", "")