llmengine/baichuanm2.py

@@ -10,22 +10,22 @@ from llmengine.base_chat_llm import BaseChatLLM, llm_register
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 class BaichuanM2LLM(BaseChatLLM):
-	def __init__(self, model_id):
+    def __init__(self, model_id):
-		self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
-		self.model = AutoModelForCausalLM.from_pretrained(
+        self.model = AutoModelForCausalLM.from_pretrained(
-			model_id,
+            model_id,
-			torch_dtype="auto",
+            torch_dtype="auto",
-			device_map="auto"
+            device_map="auto"
-		)
+        )
-		self.model_id = model_id
+        self.model_id = model_id
 
-	def build_kwargs(self, inputs, streamer):
+    def build_kwargs(self, inputs, streamer):
-		generate_kwargs = dict(
+        generate_kwargs = dict(
-			**inputs,
+            **inputs,
-			streamer=streamer,
+            streamer=streamer,
-			max_new_tokens=4096
+            max_new_tokens=4096
-		)
+        )
-		return generate_kwargs
+        return generate_kwargs
 
 	def _messages2inputs(self, messages):
 		text = tokenizer.apply_chat_template(

llmengine/base_chat_llm.py

@@ -150,14 +150,14 @@ class BaseChatLLM:
 		for d in self.output_generator(streamer):
 			if i == 0:
 				i = 1
-				t2 = time()
+				t1 = time()
 			if d['choices'][0]['finish_reason'] != 'stop':
 				txt += d['choices'][0]['delta']['content']
 			else:
 				o_tokens = d.get('output_tokens', 0)
 		i_tokens  = input_len
 
-		t3 = time()
+		t2 = time()
 		return {
 			'id': f'chatcmpl-{getID()}',
 			"object":"chat.completion",

llmengine/server.py

@@ -8,7 +8,7 @@ from llmengine.gemma3_it import Gemma3LLM
 from llmengine.medgemma3_it import MedgemmaLLM
 from llmengine.qwen3 import Qwen3LLM
 from llmengine.qwen3coder import Qwen3CoderLLM
-from llmengine.baichuanm2 import BaichuanM2LLM
+from llmengine.baiduanm2 import BaichuanM2LLM
 
 from appPublic.registerfunction import RegisterFunction
 from appPublic.log import debug, exception