BART-ender commited on
Commit
63d8572
·
verified ·
1 Parent(s): 2620860

fix(generation): handle missing chat template and add HRM-Text-1B fallback

Browse files
Files changed (1) hide show
  1. app/generation/prompting.py +19 -8
app/generation/prompting.py CHANGED
@@ -18,11 +18,22 @@ def build_messages(question: str) -> list[dict[str, str]]:
18
 
19
  def render_prompt(tokenizer: Any, question: str) -> str:
20
  messages = build_messages(question)
21
- if hasattr(tokenizer, "apply_chat_template"):
22
- return tokenizer.apply_chat_template(
23
- messages,
24
- tokenize=False,
25
- add_generation_prompt=True,
26
- )
27
-
28
- return f"{SYSTEM_PROMPT}\n\nQuestion: {question}\n"
 
 
 
 
 
 
 
 
 
 
 
 
18
 
19
  def render_prompt(tokenizer: Any, question: str) -> str:
20
  messages = build_messages(question)
21
+ if hasattr(tokenizer, "apply_chat_template") and getattr(tokenizer, "chat_template", None) is not None:
22
+ try:
23
+ return tokenizer.apply_chat_template(
24
+ messages,
25
+ tokenize=False,
26
+ add_generation_prompt=True,
27
+ )
28
+ except Exception:
29
+ pass
30
+
31
+ # Fallback for models without a chat template
32
+ # Special handling for HRM-Text-1B which requires specific condition tokens
33
+ model_name = getattr(tokenizer, "name_or_path", "")
34
+ if model_name and "HRM-Text-1B" in model_name:
35
+ # Use the synth,cot prefix as recommended for this model
36
+ prefix = "<|im_start|><|quad_end|><|object_ref_end|><|im_end|>"
37
+ return f"{prefix}{SYSTEM_PROMPT}\n\nUser: {question}\n\nAssistant:"
38
+
39
+ return f"{SYSTEM_PROMPT}\n\nUser: {question}\n\nAssistant:"