yujunjun/olmocr
1
0
Fork 0
mirror of https://github.com/allenai/olmocr.git synced 2025-10-12 08:43:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

Adjusting tagging pipelien v2

Browse Source
This commit is contained in:
Jake Poznanski 2025-05-10 17:43:56 +00:00
parent 1538163f6f
commit e06fd622c3
1 changed files with 11 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
scripts/tagging_pipeline_v2.py
View File

@ -85,6 +85,7 @@ class PIIClassification(BaseModel):
is_correspondence_or_letter: Optional[bool] is_correspondence_or_letter: Optional[bool]
is_public_order: Optional[bool] is_public_order: Optional[bool]
is_court_notice: Optional[bool] is_court_notice: Optional[bool]
is_completion_certificate: Optional[bool]
contains_pii: Optional[bool] = Field(..., description="True if document contains PII") contains_pii: Optional[bool] = Field(..., description="True if document contains PII")
@ -103,7 +104,7 @@ async def _process_single_page(page_text: str) -> PIIClassification:
"type": "text", "type": "text",
"text": ( "text": (
f"{text}\n\n-----------\n" f"{text}\n\n-----------\n"
"Given the text above, determine what type of document it is. Answer in JSON. The format of your json object should be {'primary_language': str, 'document_type': str, 'is_resume_cv': bool, 'is_academic_paper': bool, 'is_textbook': bool, 'is_news_article': bool, 'is_test_or_quiz': bool, 'is_homework_assignment': bool, 'is_class_syllabus': bool, 'is_meeting_minutes': bool, 'is_legal_contract': bool, 'is_form': bool, 'is_correspondence_or_letter': bool, 'is_public_order': bool, 'is_court_notice': bool, 'contains_pii': bool}" "Given the text above, determine what type of document it is. Answer in JSON. The format of your json object should be {'primary_language': str, 'document_type': str, 'is_resume_cv': bool, 'is_academic_paper': bool, 'is_textbook': bool, 'is_news_article': bool, 'is_test_or_quiz': bool, 'is_homework_assignment': bool, 'is_class_syllabus': bool, 'is_meeting_minutes': bool, 'is_legal_contract': bool, 'is_form': bool, 'is_correspondence_or_letter': bool, 'is_public_order': bool, 'is_court_notice': bool, 'is_completion_certificate': bool, 'contains_pii': bool}"
), ),
} }
], ],
@ -265,7 +266,9 @@ async def process_dolma_document(args, dolma_doc, sem):
# Take first 5000 characters of the document # Take first 5000 characters of the document
sample_text = text[:5000] sample_text = text[:5000]
text_length = len(text)
span_end = min(5000, text_length)
# Process the sample with the semaphore to limit concurrent requests # Process the sample with the semaphore to limit concurrent requests
async with sem: async with sem:
pii_class = await _process_single_page(sample_text) pii_class = await _process_single_page(sample_text)
@ -275,9 +278,12 @@ async def process_dolma_document(args, dolma_doc, sem):
key_name = f"{prefix}_{field_name}" key_name = f"{prefix}_{field_name}"
attribute_value = getattr(pii_class, field_name) attribute_value = getattr(pii_class, field_name)
# Add the classification result to all pages # Create a span from 0 to min(5000, len(text)) with the attribute value
for start_pos, end_pos, page_num in page_numbers: result_attributes[key_name].append([0, span_end, attribute_value])
result_attributes[key_name].append([start_pos, end_pos, attribute_value])
# If the document is longer than 5000 characters, add a null span for the rest
if text_length > 5000:
result_attributes[key_name].append([span_end, text_length, None])
return result_attributes return result_attributes
else: else: