jhansss commited on
Commit
157f247
·
1 Parent(s): 7974242

add reference link

Browse files
Files changed (1) hide show
  1. modules/llm/qwen3.py +12 -2
modules/llm/qwen3.py CHANGED
@@ -1,3 +1,5 @@
 
 
1
  from .base import AbstractLLMModel
2
  from .registry import register_llm_model
3
  from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -14,7 +16,13 @@ class Qwen3LLM(AbstractLLMModel):
14
  ).eval()
15
  self.tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=cache_dir)
16
 
17
- def generate(self, prompt: str, enable_thinking: bool = True, max_new_tokens: int = 32768, **kwargs) -> str:
 
 
 
 
 
 
18
  messages = [{"role": "user", "content": prompt}]
19
  text = self.tokenizer.apply_chat_template(
20
  messages,
@@ -23,7 +31,9 @@ class Qwen3LLM(AbstractLLMModel):
23
  enable_thinking=enable_thinking,
24
  )
25
  model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device)
26
- generated_ids = self.model.generate(**model_inputs, max_new_tokens=max_new_tokens)
 
 
27
  output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist()
28
  # parse thinking content
29
  if enable_thinking:
 
1
+ # Ref: https://qwenlm.github.io/blog/qwen3/
2
+
3
  from .base import AbstractLLMModel
4
  from .registry import register_llm_model
5
  from transformers import AutoModelForCausalLM, AutoTokenizer
 
16
  ).eval()
17
  self.tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=cache_dir)
18
 
19
+ def generate(
20
+ self,
21
+ prompt: str,
22
+ enable_thinking: bool = True,
23
+ max_new_tokens: int = 32768,
24
+ **kwargs
25
+ ) -> str:
26
  messages = [{"role": "user", "content": prompt}]
27
  text = self.tokenizer.apply_chat_template(
28
  messages,
 
31
  enable_thinking=enable_thinking,
32
  )
33
  model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device)
34
+ generated_ids = self.model.generate(
35
+ **model_inputs, max_new_tokens=max_new_tokens
36
+ )
37
  output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist()
38
  # parse thinking content
39
  if enable_thinking: