數據管護是預訓練和持續訓練模型的第一步,也可以說是最重要的一步,對于 大型語言模型 (LLM) 和小型語言模型 (SLM) 都至關重要。NVIDIA 最近宣布了 NVIDIA NeMo Curator,這是一個數據管護框架,旨在準備大規模、高質量的數據集,以用于預訓練生成式 AI 模型。
NeMo Curator 是 NVIDIA NeMo 的一部分,提供了從 Common Crawl、Wikipedia 和 arXiv 等多種公共來源下載和整理數據的開箱即用工作流。此外,它還為開發者提供了自定義數據整理流程的靈活性,以滿足其獨特需求并創建自定義數據集。
本文將介紹如何使用 NeMo Curator 創建自定義數據管護工作流。這樣一來,您可以:
- 根據生成式 AI 項目的特定需求,定制數據管護和定制流程。
- 應用嚴格的過濾器和重復數據刪除技術,使用最佳數據集訓練模型,從而確保數據質量。
- 通過識別和刪除個人識別信息 (PII) 來保護隱私,并遵守數據保護法規。
- 通過自動管理流程簡化開發流程,節省時間和資源,讓您專注于解決業務特定問題。
概述
本教程的重點是介紹如何創建一個簡單的數據管護流程,用于下載、處理和篩選 TinyStories 數據集。TinyStories 是一個由 GPT-3.5 和 GPT-4 生成的約 220 萬個短篇故事組成的數據集,其中包含 3 至 4 歲兒童能夠理解的英語單詞。該數據集在 Hugging Face 上公開提供。如果您需要詳細了解數據集,請參閱 TinyStories:語言模型的規模有多小,并且仍然會講連貫一致的英語?
此數據集體積小巧,非常適合在本地機器上創建和驗證數據管護流程。數據集分為訓練和驗證文件。本教程主要使用驗證文件,其中包含大約 22000 條記錄。
定義數據管護管道涉及以下高級步驟:
- 定義自定義文檔構建器,以便:
- 從 Web 下載數據集并轉換為 JSONL 格式。
- 對數據集進行迭代并提取每個文檔。
- 定義自定義修改器以清理和統一文本數據。
- 使用預定義和用戶定義的啟發式算法過濾數據集。
- 刪除重復數據集并刪除相同的記錄。
- 編輯數據集中的所有個人識別信息 (PII)。
- 將結果輸出為 JSONL 格式。
在消費級硬件上執行此策展流程需要不到 5 分鐘,且策展數據集在策展后應包含大約 21500 條記錄。要訪問本教程的完整代碼,請訪問 NVIDIA/NeMo-Curator 的 GitHub 倉庫。
預備知識
在開始之前,必須安裝 NeMo Curator 框架。按照項目中的說明操作,按照 NeMo Curator GitHub 自述文件 中的說明安裝框架。然后,從終端運行以下命令以驗證安裝。此外,還需要安裝后續操作所需的其他依賴項。
$ python -c "import nemo_curator; print(nemo_curator);" $ pip3 install requests |
定義自定義文檔構建器
為支持處理任意數據集,NeMo Curator 提供了一組 文檔構建器,以抽象化底層數據集的表示,包括:
DocumentDownloader
:一種用于將遠程數據下載到磁盤的抽象類。DocumentIterator
:一種用于從磁盤讀取數據集原始記錄的抽象基類。DocumentExtractor
:抽象類,用于從磁盤上的記錄中提取文本記錄及其相關元數據。
用戶可在 Omniverse 上使用多種實現,以便將其與 CommonCrawl、Wikipedia 和 arXiv 等數據集結合使用。NVIDIA/NeMo-Curator GitHub 庫提供了相關資源。以下部分將展示如何實現每個抽象類,以使用 TinyStories 數據集自定義工作流程。
下載 TinyStories 數據集
首先,實現 DocumentDownloader
類,該類負責獲取數據集驗證拆分的 URL,并使用 requests
庫。
import requests from nemo_curator.download.doc_builder import DocumentDownloader class TinyStoriesDownloader(DocumentDownloader): def __init__( self , download_dir: str ): super ().__init__() if not os.path.isdir(download_dir): os.makedirs(download_dir) self ._download_dir = download_dir print ( "Download directory: " , self ._download_dir) def download( self , url: str ) - > str : filename = os.path.basename(url) output_file = os.path.join( self ._download_dir, filename) if os.path.exists(output_file): print (f "File '{output_file}' already exists, skipping download." ) return output_file print (f "Downloading TinyStories dataset from '{url}'..." ) response = requests.get(url) with open (output_file, "wb" ) as file : file .write(response.content) return output_file |
接下來,使用以下代碼下載實際數據集:
# Download the TinyStories dataset. downloader = TinyStoriesDownloader( "/path/to/download/" ) tinystories_fp = downloader.download(TINY_STORIES_URL) write_jsonl(tinystories_fp, jsonl_dir) |
數據集將作為純文本文件下載。為解析此數據集,請使用 DocumentIterator
和 DocumentExtractor
類。這將使您能夠將其轉換為 JSONL 格式,該格式為 NeMo Curator 支持的格式之一。
從數據集中迭代和提取文本
在下載的文件中,每個記錄(或故事)跨越幾行,并通過<|endoftext|>
令牌分隔。DocumentIterator
類定義了iterate
函數,該函數會獲取要迭代的文件的路徑,并以記錄中的原始文本和(可選)該記錄的任何相關元數據的形式生成該文件的每條記錄。雖然向每條記錄添加元數據并不是強制性的,但是一些數據處理算法(例如重復數據刪除)依賴于此類數據來唯一標識每個文檔并正確執行其預期功能。
接下來,為 TinyStories 數據集實現迭代器。鑒于每個故事可以跨越多行,請定義迭代器函數,以便它會繼續讀取 (和存儲) 文件中的每行,直到到達分隔符令牌。
到達分隔符后,連接到目前為止看到的所有行,將一些元數據固定到記錄中,并生成結果。為確保記錄是唯一可識別的,請使用數據集的文件名以及內部計數器創建唯一的 id
和(可選)filename
,每條記錄中都包含這些元數據:
from nemo_curator.download.doc_builder import DocumentIterator class TinyStoriesIterator(DocumentIterator): SEPARATOR_TOKEN = "<|endoftext|>" def __init__( self ): super ().__init__() self ._counter = - 1 def iterate( self , file_path): self ._counter = - 1 file_name = os.path.basename(file_path) with open (file_path, "r" ) as file : example = [] def split_meta(example): if example: self ._counter + = 1 content = " " .join(example) meta = { "filename" : file_name, "id" : f "{file_name}-{self._counter}" , } return meta, content for line in file : if line.strip() = = TinyStoriesIterator.SEPARATOR_TOKEN: if example: yield split_meta(example) example = [] else : example.append(line.strip()) if example: yield split_meta(example) |
最后一個要實現的文檔構建器是DocumentExtractor
類,它只需返回每個記錄的文本。請注意,您可以選擇將提取的文本的一些元數據關聯起來,但這種元數據的使用超出了本教程的范圍。
from nemo_curator.download.doc_builder import DocumentExtractor class TinyStoriesExtractor(DocumentExtractor): def extract( self , content: str ) - > Tuple [ Set , str ]: # No metadata for the text, just the content. return {}, content |
將數據集寫入 JSONL 格式
NeMo Curator 提供幫助程序,可以從磁盤加載 JSONL、Parquet 或 Pickle 格式的數據集。鑒于 JSONL 格式的流行程度,本節演示如何使用之前實現的迭代器和提取器類將原始文本數據集轉換為此格式。
要將數據集轉換為 JSONL,只需將 TinyStoriesIterator
實例應用于下載的純文本文件中,對每條記錄進行迭代,并使用 TinyStoriesExtractor
實例對其進行處理。然后,從每個記錄(故事)創建一個 JSON 對象,并將其寫入輸出文件中的一行。整個過程非常簡單:
import os import json def write_jsonl(input_filename: str , output_dir: str , dump_every_n: int = 10000 ): basename = os.path.basename(input_filename) iterator = TinyStoriesIterator() extractor = TinyStoriesExtractor() to_dump = [] dump_ctr = 0 def dump_to_file(to_dump, dump_ctr): """Helper function to facilitate dumping to file.""" output_filename = f "{basename}-{dump_ctr}.jsonl" with open (os.path.join(output_dir, output_filename), "w" ) as output_file: output_file.writelines(to_dump) # Empty out the list and increment the counter. return [], dump_ctr + 1 for item in iterator.iterate(input_filename): record_meta, content = item extracted = extractor.extract(content) if extracted is None : continue text_meta, text = extracted if text is None : continue line = { "text" : text, * * text_meta, * * record_meta, } json_out = json.dumps(line, ensure_ascii = False ) to_dump.append(json_out + "\n" ) # Should we dump what we have so far? if len (to_dump) = = dump_every_n: to_dump, dump_ctr = dump_to_file(to_dump, dump_ctr) # Dump the remaining records. if to_dump: dump_to_file(to_dump, dump_ctr) |
請注意,默認情況下,此函數為每 10000 條記錄創建一個 JSONL 文件。雖然完全是可選的,但這是為了確保每個輸出文件保持足夠小,以便使用文本編輯器輕松進行手動檢查,而不會消耗過多內存。
另請注意,每個故事的內容都寫入了text
字段,該字段位于每個 JSON 對象中。整個 NeMo Curator 中的許多數據管護操作都需要知道每條記錄中的哪個字段包含該記錄的文本數據。如果未明確指定,這些操作將假設存在text
字段中的數據集。因此,通常較好的做法是始終填充每個記錄的text
字段,以及感興趣的文本數據。
使用文檔構建器加載數據集
在 NeMo Curator 中,數據集表示為類型對象 DocumentDataset
,這為從磁盤加載各種格式的數據集提供了輔助工具。例如,以 JSONL 格式創建數據集后,您可以使用以下代碼加載數據集并開始使用:
from nemo_curator.datasets import DocumentDataset # define `files` to be a list of all the JSONL files to load dataset = DocumentDataset.read_json(files, add_filename=True) |
您現在擁有了定義自定義數據集策管線和為訓練 (或驗證) 用例準備數據所需的一切。
文本清理和統一
在涉及文本數據的數據管護流程中,文本統一和清理是一項基本操作,因為從在線來源抓取的文本可能存在不一致或 Unicode 問題。為了修改文檔,NeMo Curator 提供了一個 DocumentModifier
界面,用于定義如何修改每個文檔中的給定文本。實際修改通過 Modify
幫助器實現,它需要 DocumentDataset
對象以及 DocumentModifier
,并將修飾符應用到數據集。
TinyStories 數據集具有不一致的引號,其中一些引號是卷曲的,而另一些則是筆直的。此類不一致 (例如,質量不佳的標記) 可能會給基于此數據訓練的模型帶來問題。
要解決這些問題,請創建 DocumentModifier
,它通過將所有卷曲引號替換為其直接變體,從而統一文檔中的所有單引號和雙引號:
from nemo_curator.modifiers import DocumentModifier class QuotationUnifier(DocumentModifier): def modify_document( self , text: str ) - > str : text = text.replace( "‘" , "'").replace("’", "'" ) text = text.replace( "“" , '"').replace("”", '"' ) return text |
NeMo Curator 提供了各種開箱即用的 DocumentModifier
實現。其中一個修飾符是 UnicodeReformatter
,它使用 ftfy 解決數據集中的所有 Unicode 問題。然后,將這些修改器連接在一起并清理下載的數據集。通過 Sequential
類,它接受要按順序執行的操作列表,并將這些操作應用到給定的 DocumentDataset
實例中:
from nemo_curator import Sequential from nemo_curator.modules.modify import Modify from nemo_curator.modifiers.unicode_reformatter import UnicodeReformatter def clean_and_unify(dataset: DocumentDataset) - > DocumentDataset: cleaners = Sequential( [ # Unify all the quotation marks Modify(QuotationUnifier()), # Unify all unicode Modify(UnicodeReformatter()), ] ) return cleaners(dataset) |
數據集過濾
數據集管理過程中的另一個重要步驟是數據過濾,其中會丟棄一些不符合某些標準的文檔。例如,您可能想丟棄過短、過長或不完整的文檔。在撰寫本文時,NeMo Curator 為自然語言提供了 24 種啟發式算法,為編碼語言提供了 8 種啟發式算法。
NeMo Curator 提供了一個 DocumentFilter
界面,該界面定義了一種基于各種標準對文檔進行評分的方法,并提供了一個 ScoreFilter
輔助工具,用于篩選文檔。ScoreFilter
輔助工具需要一個 DocumentDataset
和 DocumentFilter
,以確定數據集中的每個文檔是否通過篩選條件。
創建一個簡單的DocumentFilter
,用于確定故事是否以句子字符結尾。我們的目標是篩選掉所有未以句子字符結尾的故事:
from nemo_curator.filters import DocumentFilter class IncompleteStoryFilter(DocumentFilter): def __init__( self ): super ().__init__() self ._story_terminators = { "." , "!" , "?" , '"' , "”"} def score_document( self , text: str ) - > bool : ret = text.strip()[ - 1 ] in self ._story_terminators return ret def keep_document( self , score) - > bool : return score |
主要功能體現在score_document
和keep_document
函數中,其中如果文檔未以句子結束字符結尾,則返回False
(即不要保留此文檔)。
要將此過濾器應用于數據集,請傳遞一個IncompleteStoryFilter
對象作為ScoreFilter
。NeMo Curator 提供了許多開箱即用的DocumentFilter
實現。這些過濾器可以通過Sequential
類來組合使用。以下代碼展示了如何對數據集應用各種過濾器:
def filter_dataset(dataset: DocumentDataset) - > DocumentDataset: filters = Sequential( [ ScoreFilter( WordCountFilter(min_words = 80 ), text_field = "text" , score_field = "word_count" , ), ScoreFilter(IncompleteStoryFilter(), text_field = "text" ), ScoreFilter( RepeatingTopNGramsFilter(n = 2 , max_repeating_ngram_ratio = 0.2 ), text_field = "text" , ), ScoreFilter( RepeatingTopNGramsFilter(n = 3 , max_repeating_ngram_ratio = 0.18 ), text_field = "text" , ), ScoreFilter( RepeatingTopNGramsFilter(n = 4 , max_repeating_ngram_ratio = 0.16 ), text_field = "text" , ), ] ) return filters(dataset) |
此代碼將過濾所有短篇(不超過 80 個字)或不完整的故事,以及任何其他具有一定重復 n – 克比率的故事。請注意:text_field=”text”
這告訴我們:ScoreFilter
將傳遞數據集的內容:text
每個篩選條件的列。
重復數據刪除
在處理大量文本數據時,可能會存在彼此相同(或幾乎完全相同)的記錄。針對此類數據進行訓練可能會產生額外的計算和存儲開銷。NeMo Curator 提供了查找和丟棄此類重復數據的功能。為簡單起見,請專注于在數據集中查找精確的重復記錄。這可以使用 ExactDuplicates
類,如下所示。
本模塊將自動利用現有的 CUDA 設備和 GPU 加速的實現,RAPIDS cuDF 庫 識別重復文檔,從而大幅縮短處理時間。這是因為重復數據刪除階段需要計算每個文檔的哈希值,而這是計算密集型的任務。每個文檔都可以獨立進行哈希處理,因此此工作負載非常適合在 GPU 上并行運行。
from nemo_curator.modules import ExactDuplicates def dedupe(dataset: DocumentDataset) - > DocumentDataset: deduplicator = ExactDuplicates(id_field = "id" , text_field = "text" , hash_method = "md5" ) # Find the duplicates duplicates = deduplicator(dataset) docs_to_remove = duplicates.df.map_partitions( lambda x: x[x._hashes.duplicated(keep = "first" )] ) # Remove the duplicates using their IDs. duplicate_ids = list (docs_to_remove.compute(). id ) dataset_df = dataset.df deduped = dataset_df[~dataset_df. id .isin(duplicate_ids)] return DocumentDataset(deduped) |
這指定了每條記錄的唯一標識符和內容,分別是id
和text
。回想一下,在下載和提取階段,已為每個文檔分配了唯一標識符。這使得重復數據刪除器能夠以唯一方式識別彼此之間的文檔。重復數據刪除器對象將返回一組已確定為重復的 ID,只需從數據集中刪除這些文檔即可。
PII 編輯
本教程中討論的最后一個處理步驟是編輯個人識別信息 (PII)。NeMo Curator 使用 PiiModifier
類,這是對之前討論過的 DocumentModifier
類的擴展。此修飾符利用了 Presidio 框架,并允許您指定要檢測的 PII、對每次檢測采取的操作,以及批量處理數據以加速操作。
TinyStories 數據集中的故事包含許多名字的實例。此示例旨在檢測所有此類名字,并將其替換為匿名令牌。這可以使用幾行代碼來完成:
from nemo_curator.modifiers.pii_modifier import PiiModifier def redact_pii(dataset: DocumentDataset) - > DocumentDataset: redactor = Modify( PiiModifier( supported_entities = [ "PERSON" ], anonymize_action = "replace" , device = "cpu" , ), ) return redactor(dataset) |
此操作會獲取整個數據集,并返回修改后的數據集。
整合管線
在實現管線的每個步驟后,是時候將所有內容放在一起,并按順序對數據集應用每個操作了。您可以使用Sequential
API 將類鏈式管理操作結合在一起:
curation_steps = Sequential( [ clean_and_unify, filter_dataset, dedupe, redact_pii, ] ) dataset = curation_steps(dataset) print ( "Executing the pipeline..." ) dataset = dataset.persist() dataset.to_json( "/output/path" , write_to_filename = True ) |
在幕后,NeMo Curator 使用 Dask 以分布式方式處理數據集。由于 Dask 操作是延遲評估的,因此有必要調用.persist
方法,以指示 Dask 應用操作的函數。處理完成后,您可以通過調用.to_json
方法,并提供輸出路徑,以獲取處理結果。
后續步驟
NeMo Curator 支持許多先進的數據處理和過濾技術,例如模糊或基于任務的重復數據刪除、任務識別和去污染、域分類等(還有更多內容),本教程中并未涵蓋這些技術。欲了解更多信息,請查看 GitHub 上的數據管護示例集合。
此外,您還可以 申請訪問 NVIDIA NeMo Curator 微服務,該微服務為企業提供了從任何地方開始數據管護的最簡單途徑。它提供了精簡的性能和可擴展性,以縮短上市時間。
?