Source code for oumi.datasets.vision_language.geometry3k
# Copyright 2025 - Oumi## Licensed under the Apache License, Version 2.0 (the "License");# you may not use this file except in compliance with the License.# You may obtain a copy of the License at## http://www.apache.org/licenses/LICENSE-2.0## Unless required by applicable law or agreed to in writing, software# distributed under the License is distributed on an "AS IS" BASIS,# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.# See the License for the specific language governing permissions and# limitations under the License.fromtyping_extensionsimportoverridefromoumi.core.datasetsimportVisionLanguageSftDatasetfromoumi.core.registryimportregister_datasetfromoumi.core.types.conversationimport(ContentItem,Conversation,Message,Role,Type,)_GEOMETRY3K_INSTRUCTION=(r"You FIRST think about the reasoning process as an internal monologue and then "r"provide the final answer. "r"The reasoning process MUST BE enclosed within <think> </think> tags. "r"The final answer MUST BE put in \boxed{}.")
[docs]@register_dataset("hiyouga/geometry3k")classGeometry3kDataset(VisionLanguageSftDataset):"""Dataset class for the `hiyouga/geometry3k` dataset."""default_dataset="hiyouga/geometry3k"def__init__(self,*,add_system_instruction:bool=False,**kwargs,)->None:"""Initializes a new instance of the Geometry3kDataset class."""self._add_system_instruction=add_system_instructionsuper().__init__(**kwargs)def_process_text_value(self,s:str)->str:# The data contains occasional `\n` at the beginning or end# of text values. Let's strip them.returns.strip()ifselse""
[docs]@overridedeftransform_conversation(self,example:dict)->Conversation:"""Transform a single conversation example into a Conversation object."""input_text=self._process_text_value(example["problem"])input_text=input_text.removeprefix("<image>")ifnotself._add_system_instruction:input_text=input_text+" "+_GEOMETRY3K_INSTRUCTIONoutput_text=self._process_text_value(example["answer"])image=example["images"]iflen(image)!=1:raiseValueError(f"Expected 1 image, but got {len(image)} images in example.")image=image[0]messages=([Message(role=Role.SYSTEM,content=_GEOMETRY3K_INSTRUCTION,)]ifself._add_system_instructionelse[])+[Message(role=Role.USER,content=[ContentItem(type=Type.IMAGE_BINARY,binary=image["bytes"],),ContentItem(type=Type.TEXT,content=input_text),],),Message(role=Role.ASSISTANT,content=output_text),]returnConversation(messages=messages)