File size: 2,369 Bytes
b2e30fd |
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 |
import os
from transformers import DataCollatorForLanguageModeling, Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM
from datasets import load_dataset, Dataset
os.environ['TOKENIZERS_PARALLELISM'] = 'true'
def load_transcripts(directory):
transcript_texts = []
for filename in os.listdir(directory):
if filename.endswith('.txt'):
file_path = os.path.join(directory, filename)
with open(file_path, 'r', encoding='utf-8') as file:
transcript_texts.append(file.read())
return transcript_texts
def main():
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
tokenizer.pad_token = tokenizer.eos_token # Set a padding token
datasets = load_dataset(
'csv', data_files={'test': 'test.csv'})
test_dataset = datasets['test'].map(
lambda examples: tokenizer(
examples['transcript'], truncation=True, padding='max_length', return_attention_mask=False),
batched=True
)
transcripts = load_transcripts('transcripts')
train_encodings = tokenizer(transcripts, truncation=True,
padding='max_length', max_length=tokenizer.model_max_length)
train_dataset = Dataset.from_dict(
{"input_ids": train_encodings["input_ids"], "attention_mask": train_encodings["attention_mask"]})
data_collator = DataCollatorForLanguageModeling(
tokenizer=tokenizer, mlm=False)
training_args = TrainingArguments(
output_dir="./model",
overwrite_output_dir=True,
num_train_epochs=3,
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
save_steps=2_000,
save_total_limit=2,
logging_dir='./logs',
learning_rate=5e-5,
weight_decay=0.01,
evaluation_strategy="steps",
eval_steps=500,
warmup_steps=300
)
trainer = Trainer(
model=model,
args=training_args,
data_collator=data_collator,
train_dataset=train_dataset,
eval_dataset=test_dataset,
)
# Training
trainer.train()
trainer.evaluate()
# Saving the trained model
model.save_pretrained("./model")
tokenizer.save_pretrained("./model")
if __name__ == "__main__":
main()
|