🔨[DEV] Add model generation logger

src/fastmindapi/model/llama_cpp/LLM.py

@@ -1,14 +1,14 @@
 from typing import Optional
 
+from ..utils.io import generation_logger
 from ...utils.transform import convert_numpy_float32_to_float, clean_dict_null_value
-from ... import logger
-from ... import config as fmconfig
 
 class LlamacppLLM:
     def __init__(self, 
                  model):
         self.model = model
         self.model_name = None
+        self.backend = "Llamacpp"
 
     @classmethod
     def from_path(cls, 
@@ -24,6 +24,7 @@ def __call__(self,
         return output_text
         # {"id":"cmpl-bab2b133-cf08-43aa-8ea0-7c4b109b9cf4","object":"text_completion","created":1726721257,"model":"/Users/wumengsong/Resource/gguf/Meta-Llama-3.1-8B-Instruct-Q8_0.gguf","choices":[{"text":" I'm a beginner and I'ts my first time playing this game. I","index":0,"logprobs":null,"finish_reason":"length"}],"usage":{"prompt_tokens":9,"completion_tokens":16,"total_tokens":25}}
 
+    @generation_logger
     def generate(self,
                  input_text: str,
                  max_new_tokens: Optional[int] = None,
@@ -73,10 +74,6 @@ def generate(self,
                     logits["probs"].append(round(math.exp(logprobs["top_logprobs"][i][token]),4))
                 logits_list.append(logits)
 
-        if fmconfig.log_model_io:
-            logger.info("【model_io】Llamacpp:"+self.model_name+".generate()")
-            logger.info("- input_text: "+input_text)
-            logger.info("- output_text: "+output_text)
         generation_output = {"output_text": output_text,
                              "input_id_list": input_id_list,
                              "input_token_list": input_token_list,

src/fastmindapi/model/openai/ChatModel.py

@@ -1,8 +1,10 @@
 from typing import Optional
 
-from ...utils.transform import convert_openai_logprobs, clean_dict_null_value
 from ... import logger
-from ... import config as fmconfig
+from ..utils.io import generation_logger
+from ..utils.transform import convert_openai_logprobs
+from ...utils.transform import clean_dict_null_value
+
 
 class OpenAIChatModel:
     def __init__(self, 
@@ -12,6 +14,7 @@ def __init__(self,
         self.client = client
         self.system_prompt = system_prompt
         self.model_name = model_name
+        self.backend = "OpenAI"
 
     @classmethod
     def from_client(cls, 
@@ -34,6 +37,7 @@ def __call__(self, input_text: str,
         except Exception as e:
             return "【Error】: " + str(e)
 
+    @generation_logger
     def generate(self,
                  input_text: str,
                  max_new_tokens: Optional[int] = None,
@@ -66,10 +70,6 @@ def generate(self,
         logits_list = None
         if return_logits:
             logits_list = convert_openai_logprobs(completion.choices[0].logprobs)
-        if fmconfig.log_model_io:
-            logger.info("【model_io】OpenAI:"+self.model_name+".generate()")
-            logger.info("- input_text: "+input_text)
-            logger.info("- output_text: "+output_text)
         generation_output = {"output_text": output_text,
                             #  "input_id_list": input_id_list,
                             #  "input_token_list": input_token_list,

src/fastmindapi/model/transformers/CausalLM.py

@@ -1,7 +1,6 @@
 from typing import Optional
 
-from ... import logger
-from ... import config as fmconfig
+from ..utils.io import generation_logger
 from ...utils.transform import clean_dict_null_value
 
 class TransformersCausalLM:
@@ -11,7 +10,7 @@ def __init__(self,
         self.tokenizer = tokenizer
         self.model = model
         self.model_name = None
-
+        self.backend = "Transformers"
         self.model.eval()
 
 
@@ -35,6 +34,7 @@ def __call__(self,
         output_text = full_text[len(re_inputs):]
         return output_text
 
+    @generation_logger
     def generate(self,
                  input_text: str,
                  max_new_tokens: Optional[int] = None,
@@ -111,10 +111,6 @@ def generate(self,
                     logits["probs"].append(round(prob,4))
                 logits_list.append(logits)
 
-        if fmconfig.log_model_io:
-            logger.info("【model_io】Transformers:"+self.model_name+".generate()")
-            logger.info("- input_text: "+input_text)
-            logger.info("- output_text: "+output_text)
         generation_output = {"output_text": output_text,
                              "input_id_list": input_id_list,
                              "input_token_list": input_token_list,

src/fastmindapi/model/utils/io.py

@@ -0,0 +1,15 @@
+import functools
+from typing import Callable, Any
+
+from ... import config, logger
+
+def generation_logger(func: Callable) -> Callable:
+    @functools.wraps(func)  # 保留 func 的元信息
+    def wrapper(self, *args, **kwargs) -> Any:
+        result = func(self, *args, **kwargs)
+        if config.log_model_io:
+            logger.info("【model_io】"+self.backend+":"+self.model_name+".generate()")
+            logger.info("- input_text: "+result["input_text"])
+            logger.info("- output_text: "+result["output_text"])
+        return result
+    return wrapper

src/fastmindapi/model/utils/transform.py

@@ -0,0 +1,19 @@
+import math
+
+def convert_openai_logprobs(logprobs):
+    logprobs = logprobs.model_dump()
+    logits_list = []
+    for token_info in logprobs["content"]:
+        logits = {
+            "token": token_info["token"],
+            "pred_token": [],
+            # "logits": [],
+            "probs": [],
+            "logprobs": []
+        }
+        for predict_info in token_info["top_logprobs"]:
+            logits["pred_token"].append(predict_info["token"])
+            logits["logprobs"].append(round(predict_info["logprob"],4))
+            logits["probs"].append(round(math.exp(predict_info["logprob"]),4))
+        logits_list.append(logits)
+    return logits_list

src/fastmindapi/model/vllm/LLM.py

@@ -1,7 +1,6 @@
 from typing import Optional
 
-from ... import logger
-from ... import config as fmconfig
+from ..utils.io import generation_logger
 from ...utils.transform import clean_dict_null_value
 
 class vLLMLLM:
@@ -10,6 +9,7 @@ def __init__(self,
         self.model = model
         self.tokenizer = self.model.get_tokenizer()
         self.model_name = None
+        self.backend = "vLLM"
 
     @classmethod
     def from_path(cls, 
@@ -25,6 +25,7 @@ def __call__(self,
         output_text = outputs[0].outputs[0].text
         return output_text
 
+    @generation_logger
     def generate(self,
                  input_text: str,
                  max_new_tokens: Optional[int] = None,
@@ -89,10 +90,6 @@ def generate(self,
                         logits["logprobs"][rank-1] = round(logprob,4) if logprob != float("-inf") else None
                 logits_list.append(logits)
 
-        if fmconfig.log_model_io:
-            logger.info("【model_io】vLLM:"+self.model_name+".generate()")
-            logger.info("- input_text: "+input_text)
-            logger.info("- output_text: "+output_text)
         generation_output = {"output_text": output_text,
                              "input_id_list": input_id_list,
                              "input_token_list": input_token_list,

src/fastmindapi/utils/transform.py

@@ -21,21 +21,3 @@ def clean_dict_null_value(d):
             continue
         new_d[k] = d[k]
     return new_d
-
-def convert_openai_logprobs(logprobs):
-    logprobs = logprobs.model_dump()
-    logits_list = []
-    for token_info in logprobs["content"]:
-        logits = {
-            "token": token_info["token"],
-            "pred_token": [],
-            # "logits": [],
-            "probs": [],
-            "logprobs": []
-        }
-        for predict_info in token_info["top_logprobs"]:
-            logits["pred_token"].append(predict_info["token"])
-            logits["logprobs"].append(round(predict_info["logprob"],4))
-            logits["probs"].append(round(math.exp(predict_info["logprob"]),4))
-        logits_list.append(logits)
-    return logits_list