|
import logging |
|
from transformers import T5Tokenizer, T5ForConditionalGeneration, Trainer, TrainingArguments |
|
from datasets import Dataset |
|
from sklearn.model_selection import train_test_split |
|
import re |
|
from transformers import T5Tokenizer, T5ForConditionalGeneration |
|
|
|
model_name = "t5-base" |
|
tokenizer = T5Tokenizer.from_pretrained(model_name) |
|
model = T5ForConditionalGeneration.from_pretrained(model_name) |
|
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") |
|
logger = logging.getLogger(__name__) |
|
|
|
stop_words = {"and", "or", "but", "the", "is", "are", "was", "were", "a", "an", "in", "on", "at", "of", "to", "with"} |
|
def stem_word(word): |
|
suffixes = ['ing', 'ed', 'ly', 's', 'es', 'er'] |
|
for suffix in suffixes: |
|
if word.endswith(suffix): |
|
return word[:-len(suffix)] |
|
return word |
|
|
|
def clean_text(text): |
|
text = re.sub(r'[^\w\s]', '', text) |
|
text = re.sub(r'\d+', '', text) |
|
text = text.lower() |
|
text = " ".join([word for word in text.split() if word not in stop_words]) |
|
text = " ".join([stem_word(word) for word in text.split()]) |
|
return text |
|
|
|
def read_prompts(file_path): |
|
input_texts = [] |
|
target_texts = [] |
|
with open(file_path, "r", encoding="utf-8") as file: |
|
lines = file.readlines() |
|
for line in lines: |
|
if line.startswith("input:"): |
|
input_texts.append(line.replace("input:", "").strip()) |
|
elif line.startswith("target:"): |
|
target_texts.append(line.replace("target:", "").strip()) |
|
return input_texts, target_texts |
|
|
|
def prepare_data(input_texts, target_texts): |
|
inputs = tokenizer(input_texts, max_length=512, truncation=True, padding="max_length") |
|
targets = tokenizer(target_texts, max_length=512, truncation=True, padding="max_length") |
|
return {"input_ids": inputs["input_ids"], "attention_mask": inputs["attention_mask"], "labels": targets["input_ids"]} |
|
|
|
def fine_tune_model(): |
|
model_name = "./fine_tuned_model" |
|
tokenizer = T5Tokenizer.from_pretrained(model_name) |
|
model = T5ForConditionalGeneration.from_pretrained(model_name) |
|
|
|
try: |
|
logger.info("Reading and cleaning prompts.") |
|
input_texts, target_texts = read_prompts("prompts.txt") |
|
input_texts_cleaned = [clean_text(text) for text in input_texts] |
|
target_texts_cleaned = [clean_text(text) for text in target_texts] |
|
|
|
logger.info("Splitting dataset into training and validation sets.") |
|
train_texts, val_texts, train_labels, val_labels = train_test_split(input_texts_cleaned, target_texts_cleaned, test_size=0.1) |
|
|
|
logger.info("Preparing datasets for training.") |
|
train_dataset = Dataset.from_dict(prepare_data(train_texts, train_labels, tokenizer)) |
|
val_dataset = Dataset.from_dict(prepare_data(val_texts, val_labels, tokenizer)) |
|
|
|
training_args = TrainingArguments( |
|
output_dir="./results", |
|
evaluation_strategy="steps", |
|
learning_rate=5e-5, |
|
per_device_train_batch_size=4, |
|
num_train_epochs=3, |
|
save_steps=500, |
|
logging_dir="./logs", |
|
logging_steps=10 |
|
) |
|
|
|
logger.info("Starting model training.") |
|
trainer = Trainer( |
|
model=model, |
|
args=training_args, |
|
train_dataset=train_dataset, |
|
eval_dataset=val_dataset |
|
) |
|
trainer.train() |
|
|
|
logger.info("Saving fine-tuned model.") |
|
model.save_pretrained("./fine_tuned_model") |
|
tokenizer.save_pretrained("./fine_tuned_model") |
|
|
|
except Exception as e: |
|
logger.error(f"An error occurred during fine-tuning: {str(e)}") |
|
|
|
fine_tune_model() |
