Cut HF completions before EOS so the last grid row is parsed

#2
Files changed (1) hide show
  1. marc.py +4 -3
marc.py CHANGED
@@ -560,9 +560,10 @@ class HFModel:
560
  kwargs["temperature"] = temperature
561
  out = self.model.generate(**data, **kwargs)
562
  sequences = out[:, width:].tolist()
563
- # Preserve EOS and <think> spellings as before; remove only batch padding
564
- # after the first EOS so batching does not change the parser's input.
565
- sequences = [s[:next((i + 1 for i, t in enumerate(s) if t in eos), len(s))]
 
566
  for s in sequences]
567
  decoded = self.tok.batch_decode(sequences, skip_special_tokens=False)
568
  if len(decoded) != len(indices) * n:
 
560
  kwargs["temperature"] = temperature
561
  out = self.model.generate(**data, **kwargs)
562
  sequences = out[:, width:].tolist()
563
+ # Cut each sequence before its first EOS (dropping batch padding too). A kept
564
+ # EOS is decoded onto the last grid row ("0 1<|im_end|>"), which then fails
565
+ # _ROW and loses that row. <think> spellings are preserved.
566
+ sequences = [s[:next((i for i, t in enumerate(s) if t in eos), len(s))]
567
  for s in sequences]
568
  decoded = self.tok.batch_decode(sequences, skip_special_tokens=False)
569
  if len(decoded) != len(indices) * n: