# 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.fromtypingimportUnionimportpandasaspdfromoumi.core.datasetsimportBaseSftDatasetfromoumi.core.registryimportregister_datasetfromoumi.core.types.conversationimportConversation,Message,Role
[docs]@register_dataset("CohereForAI/aya_dataset")classAyaDataset(BaseSftDataset):"""Dataset class for the CohereForAI/aya_dataset dataset."""default_dataset="CohereForAI/aya_dataset"
[docs]deftransform_conversation(self,example:Union[dict,pd.Series])->Conversation:"""Transform a dataset example into a Conversation object."""question:str=example.get("inputs",None)or""answer:str=example.get("targets",None)or""messages=[Message(role=Role.USER,content=question),Message(role=Role.ASSISTANT,content=answer),]returnConversation(messages=messages)