truncate after bos

app.py

@@ -17,6 +17,13 @@ def generate_response(prompt):
     outputs = model.generate(**inputs, max_new_tokens=5, temperature=1.0)
     input_length = inputs['input_ids'].shape[1]
     new_token_ids = outputs[0][input_length:]
+    bos_token_id = tokenizer.bos_token_id
+    if bos_token_id is not None:
+        bos_positions = (new_token_ids == bos_token_id).nonzero(as_tuple=True)[0]
+        if len(bos_positions) > 0:
+            # Truncate at first BOS token
+            first_bos_pos = bos_positions[0].item()
+            new_token_ids = new_token_ids[:first_bos_pos]
     new_tokens = tokenizer.decode(new_token_ids, skip_special_tokens=False)
     return new_tokens