Skip to content

Commit 9ec8d73

Browse files
fix: add config for enabling thinking
1 parent 04a9744 commit 9ec8d73

1 file changed

Lines changed: 4 additions & 4 deletions

File tree

graphgen/models/llm/local/vllm_wrapper.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,8 @@ def _build_inputs(self, prompt: str, history: Optional[List[dict]] = None) -> An
5151
return self.tokenizer.apply_chat_template(
5252
messages,
5353
tokenize=False,
54-
add_generation_prompt=True
54+
add_generation_prompt=True,
55+
enable_thinking=False
5556
)
5657

5758
async def _consume_generator(self, generator):
@@ -72,7 +73,7 @@ async def generate_answer(
7273
temperature=self.temperature if self.temperature >= 0 else 1.0,
7374
top_p=self.top_p if self.top_p >= 0 else 1.0,
7475
max_tokens=extra.get("max_new_tokens", 2048),
75-
repetition_penalty=extra.get("repetition_penalty", 1.05),
76+
repetition_penalty=extra.get("repetition_penalty", 1.05)
7677
)
7778

7879
try:
@@ -101,8 +102,7 @@ async def generate_topk_per_token(
101102
sp = self.SamplingParams(
102103
temperature=0,
103104
max_tokens=1,
104-
logprobs=self.top_k,
105-
chat_template_kwargs={"enable_thinking": False}
105+
logprobs=self.top_k
106106
)
107107

108108
try:

0 commit comments

Comments
 (0)