mirrors/docling
1
0
Fork 0
mirror of https://github.com/DS4SD/docling.git synced 2025-12-10 13:48:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

More elegant solution in removing the input prompt

Browse Source 
Signed-off-by: Maksym Lysak <mly@zurich.ibm.com>
This commit is contained in:
Maksym Lysak
2025-02-12 18:48:48 +01:00
parent e486eb1720
commit 66532eadb6
1 changed files with 3 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
docling/models/smol_docling_model.py
View File

@@ -117,17 +117,12 @@ class SmolDoclingModel(BasePageModel):
) )
generation_time = time.time() - start_time generation_time = time.time() - start_time
generated_texts = self.processor.batch_decode( generated_texts = self.processor.batch_decode(
generated_ids, skip_special_tokens=False generated_ids[:, inputs["input_ids"].shape[1] :],
skip_special_tokens=False,
)[0] )[0]
num_tokens = len(generated_ids[0]) num_tokens = len(generated_ids[0])
# DELETE NOISE BEFORE "Assistant: "
starting_point = "Assistant: "
generated_texts = generated_texts[
generated_texts.index(starting_point) + len(starting_point) :
]
# generated_texts = generated_texts.replace("Assistant: ", "")
page_tags = generated_texts page_tags = generated_texts
inference_time = time.time() - start_time inference_time = time.time() - start_time