File size: 4,616 Bytes
f05550b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
'''
File Name: prepare_sft_dataset.py Author: Nikhil Malhotra
Date: 21/7/2024
purpose: The purpose of a file is to create high quality SFT dataset for Project Indus.
Dataset source is obtained from Hugging face and enables to get high quality SFT dataset
Dataset is then translated in requisite dialects as supported by Google
Dataset is also split in train and test and enables to create requisite files
Name of the file carries the source, translated into a dialect along with split type
'''
#Imports
import os
import pandas as pd
import configparser
from googletrans import Translator
#Configuration settings
source_dir_path = "source"
translated_dir_path = "translated"
config = configparser.ConfigParser(default_section="DATASETS")
'''
Dialects double codes used for translation
1. English : en
2. Hindi : hi
3. Dogri : doi
4. Bhojpuri: bho
5. Maithili :mai
'''
translated_dialect = "hi"
def store_sft_dataset(name_of_dataset , data_frame,split_type):
"""
Method to get data from Hugging face source and store this data in csv file
This method also calls translation in dialects
Input : name of dataset to be fetched from hugging face
Data Frame : Data frame contains question or answers or prompts and messages
Split type: test or train as the data is split
"""
file_name = str(os.path.join(source_dir_path,name_of_dataset.replace("/","_")))+"train.csv"
if not os.path.isfile(file_name):
print("Opening file....." , file_name)
data_frame.to_csv(file_name, encoding="utf-8",index=False,header=True)
print("Finished writing file....",file_name)
#Change dialect to translate at top depending upon the dialect needed
translate_in_dialects(name_of_dataset , data_frame , split_type, translated_dialect)
def translate_in_dialects(name_of_dataset , data_frame ,split_type, dialect_name="hi"):
USE_COUNT = False #Only use this if you need to test with small data
count = 0
"""
This method enables translation in various dialects byt taking a dataset in a data frame and then
storing it as a csv file
Input : name of dataset to be fetched from hugging face
Data Frame : Data frame contains question or answers or prompts and messages
Split type: test or train as the data is split
dialect name: Dialect in whcih data needs to be converted into
"""
print("Translating now....")
translator = Translator()
#Make a new dataframe for translation
translate_df = pd.DataFrame(columns=["question" , "answer"])
translated_append_list = []
for index,val in data_frame.iterrows():
translated_ques = translator.translate(val["question"], dest=dialect_name).text
translated_ans = translator.translate(val["answer"], dest=dialect_name).text
translated_append_list.append({'question' : translated_ques, 'answer': translated_ans})
count +=1
#If count is used then the program thinks you are testing.
# It enables you to break the loop with 5 data points
if USE_COUNT:
if count == 5:
break
df = pd.concat([translate_df, pd.DataFrame(translated_append_list)])
translated_file_name = str(os.path.join(translated_dir_path,
name_of_dataset.replace("/","_")
))+split_type+"_"+dialect_name+"_translated.csv"
if not os.path.isfile(translated_file_name):
print("Opening file....." , translated_file_name)
df.to_csv(translated_file_name, encoding="utf-8",index=False,header=True)
print("Finished writing file....",translated_file_name)
if __name__ == "__main__":
"""
Main method to read config file and then use it to store and translate to produce
a high quality SFT dataset
"""
#Read the config file and get the requisite section to call for the dataset
config.read("dataset.ini", encoding="utf-8")
for key in config['openai']:
if key.lower().strip().__eq__("name"):
name_of_dataset = config['openai'][key].replace('"','')
splits = { 'train': 'main/train-00000-of-00001.parquet',
'test': 'main/test-00000-of-00001.parquet'
}
df_1 = pd.read_parquet("hf://datasets/"+name_of_dataset + splits["test"])
store_sft_dataset(name_of_dataset , df_1,"train")
|