Upload finetune.py

finetune.py

@@ -0,0 +1,96 @@
+
+
+from datasets import load_dataset
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+import torch, csv
+file_dict = {
+  "train" : "name_dataset.csv",
+  "test" : "name_dataset.csv"
+}
+
+dataset = load_dataset(
+  'csv',
+  data_files=file_dict,
+  delimiter=',',
+  column_names=['text', 'label'],
+  skiprows=1
+)
+
+print(f"Train dataset size: {len(dataset['train'])}")
+print(f"Test dataset size: {len(dataset['test'])}")
+
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from datasets import concatenate_datasets
+
+model_id = "t5-small"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+def tokenize_function(example):
+    model_inputs = tokenizer(example["text"], truncation=True)
+    targets = tokenizer(example["label"], truncation=True)
+    model_inputs['labels'] = targets['input_ids']
+    return model_inputs
+
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
+tokenized_datasets = tokenized_datasets.remove_columns("text")
+tokenized_datasets = tokenized_datasets.remove_columns("label")
+
+from transformers import DataCollatorForSeq2Seq
+model =T5ForConditionalGeneration.from_pretrained(model_id)
+
+from peft import LoraConfig, get_peft_model,TaskType
+
+
+lora_config = LoraConfig(
+ r=16, 
+ lora_alpha=32,
+ target_modules=["q", "v"],
+ lora_dropout=0.05,
+ bias="none",
+ task_type=TaskType.SEQ_2_SEQ_LM
+)
+model = get_peft_model(model, lora_config)
+model.print_trainable_parameters()
+
+label_pad_token_id = -100
+data_collator = DataCollatorForSeq2Seq(
+    tokenizer,
+    model=model,
+    label_pad_token_id=label_pad_token_id,
+    pad_to_multiple_of=8
+)
+
+from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments
+
+output_dir = "lora-t5"
+training_args = Seq2SeqTrainingArguments(
+    output_dir=output_dir,
+    auto_find_batch_size=True,
+    learning_rate=1e-3,
+    num_train_epochs=100,
+    logging_dir=f"{output_dir}/logs",
+    logging_strategy="steps",
+    logging_steps=500,
+    save_strategy="no",
+#    report_to="tensorboard",
+)
+
+trainer = Seq2SeqTrainer(
+    model=model,
+    args=training_args,
+    data_collator=data_collator,
+    train_dataset=tokenized_datasets["train"],
+)
+model.config.use_cache = False
+trainer.train()
+peft_model_id = "name-peft"
+trainer.model.save_pretrained(peft_model_id)
+tokenizer.save_pretrained(peft_model_id)
+from transformers import T5ForConditionalGeneration, AutoTokenizer
+from peft import PeftModel
+base_model = T5ForConditionalGeneration.from_pretrained(model_id)
+peft_model = PeftModel.from_pretrained(base_model, "name-peft")
+peft_model = peft_model.merge_and_unload()
+peft_model.save_pretrained("name-extraction")
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+tokenizer.save_pretrained("name-extraction")