TrainModelSeq2SeqSummary.py

import torch
from transformers import BartTokenizer, BartForConditionalGeneration, pipeline
from torch.utils.data import DataLoader
import json
from tqdm import tqdm

def train_model(model, dataloader, optimizer, epochs): # 'cpu' to 'cuda' for GPU
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        total_eval = 0
        for batch in tqdm(dataloader, desc=f"Training epoch {epoch+1}/{epochs}"):
            inputs = tokenizer(batch['input_text'], return_tensors='pt', truncation=True, padding=True)
            outputs = tokenizer(batch['output_text'], return_tensors='pt', truncation=True, padding=True)
            loss = model(input_ids=inputs['input_ids'].to('cpu'), labels=outputs['input_ids'].to('cpu')).loss
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            optimizer.zero_grad()
            total_loss += loss.item()
            total_eval += 1
        avg_train_loss = total_loss / total_eval
        print(f"Epoch {epoch+1}/{epochs} - Average training loss: {avg_train_loss:.2f}")

def TrainCpu(dataset_path, pretrained_model_path, save_model_path):
    global tokenizer
    tokenizer = BartTokenizer.from_pretrained(pretrained_model_path)
    model = BartForConditionalGeneration.from_pretrained(pretrained_model_path)
    model = model.to('cpu')
    optimizer = torch.optim.Adam(model.parameters(), lr=0.000001)

    with open(dataset_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    data = [{'input_text': input_text, 'output_text': label} for input_text, label in zip(data['input_ids'], data['labels'])]
    dataloader = DataLoader(data, batch_size=1, shuffle=True)
    train_model(model, dataloader, optimizer, epochs=30)
    model.save_pretrained(save_model_path)
    print(f"Model has been trained and saved to '{save_model_path}'")

def interference(model_path, pretrained_model_path):
    tokenizer = BartTokenizer.from_pretrained(pretrained_model_path)
    model = BartForConditionalGeneration.from_pretrained(model_path)
    summarizer = pipeline('summarization', model=model, tokenizer=tokenizer)

    #text = "pristine quality, White hair, detailed, bright green eyes, breezy, flowing hair, sunny, upper body, detailed face, summer, lush greenery, golden sunlight" # summarizeExample1 = "White hair, detailed bright green eyes, summer"
    #text = "best quality, masterpiece, White hair,detailed, red eyes, windy, floating hair, snowy, upper body, detailed face, winter, trees, sunshine"
    text = "The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in Paris, France. It is named after the engineer Gustave Eiffel, whose company designed and built the tower."

    summary = summarizer(text, max_length=20, min_length=5, do_sample=False)
    print(summary[0]['summary_text'])

def main():
    # TrainCpu('dataset_CLIP.json', 'sshleifer/distilbart-cnn-12-6', 'trained_model')
    interference('trained_model', 'sshleifer/distilbart-cnn-12-6')


if __name__ == "__main__":
    main()