View file src/colab/llm_training.py - Download

# -*- coding: utf-8 -*-
"""llm_training.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/158KviNmbzhwHI3y1Lb131YSYPC0cYp3M

Links :

https://medium.com/versent-tech-blog/training-an-llm-with-hugging-face-363942f1a237

https://colab.research.google.com/drive/1VGWi8bwKixinanSW-9nkt15Ql8r4ml1g

https://colab.research.google.com/drive/1qooYZkERg7zPMdU93cm7EfxALtDxbAHP#scrollTo=bf1P8G8P_1NH

Imports
"""

from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import Trainer, Seq2SeqTrainer, AutoModelForSeq2SeqLM
from transformers import Seq2SeqTrainingArguments
from peft import LoraConfig, TaskType, get_peft_model
from transformers import pipeline

"""Questions and answers for additional training"""

# Commented out IPython magic to ensure Python compatibility.
# %%writefile acdc_qa.json
# {"question":"When was AC/DC formed?", "answer": "1973"}
# {"question":"What year was AC/DC formed?", "answer": "1973"}
# {"question":"What is the name of the band that was formed in Sydney in 1973?", "answer": "AC/DC"}
# {"question":"Where was AC/DC formed?", "answer": "Sydney"}
# {"question":"Who are the members of ACDC?", "answer": "Malcolm Young and Angus Young"}
# {"question":"What role does Malcolm play in AC/DC?", "answer": "rhythm guitarist"}
# {"question":"What role does Angus play in AC/DC?", "answer": "lead guitarist"}{"question":"What genre of music is AC/DC known for?", "answer":"Hard rock"}
# {"question":"Who was the original lead vocalist of AC/DC?", "answer":"Dave Evans"}
# {"question":"Who replaced Dave Evans as AC/DC's singer in 1974?", "answer":"Bon Scott"}
# {"question":"Which singer joined AC/DC after Bon Scott's death?", "answer":"Brian Johnson"}
# {"question":"When did Bon Scott die?", "answer":"1980"}
# {"question":"What is the title of AC/DC's first album?", "answer":"High Voltage"}
# {"question":"In what year was the album High Voltage released?", "answer":"1975"}
# {"question":"What is AC/DC's best-selling album?", "answer":"Back in Black"}
# {"question":"In what year was Back in Black released?", "answer":"1980"}
# {"question":"Who produced the album Back in Black?", "answer":"Mutt Lange"}
# {"question":"What is the name of AC/DC's logo font style?", "answer":"A variation of the Squealer typeface"}
# {"question":"What does the band's name AC/DC stand for?", "answer":"Alternating Current/Direct Current"}
# {"question":"Who founded AC/DC?", "answer":"Malcolm Young and Angus Young"}
# {"question":"Which AC/DC guitarist is known for wearing a schoolboy uniform on stage?", "answer":"Angus Young"}
# {"question":"What instrument does Angus Young play?", "answer":"Lead guitar"}
# {"question":"What instrument did Malcolm Young play?", "answer":"Rhythm guitar"}
# {"question":"What is the title of AC/DC's first international album release?", "answer":"High Voltage"}
# {"question":"What is the name of AC/DC's album released in 1979?", "answer":"Highway to Hell"}
# {"question":"Who produced the album Highway to Hell?", "answer":"Mutt Lange"}
# {"question":"What year was the album Highway to Hell released?", "answer":"1979"}
# {"question":"Which AC/DC member is Malcolm and Angus Young's older brother?", "answer":"George Young"}
# {"question":"What role did George Young have with AC/DC?", "answer":"Producer and mentor"}
# {"question":"Who played bass for AC/DC for the longest period?", "answer":"Cliff Williams"}
# {"question":"When did Cliff Williams join AC/DC?", "answer":"1977"}
# {"question":"Who was AC/DC's first drummer?", "answer":"Colin Burgess"}
# {"question":"Which AC/DC drummer was known for long tenures in the band?", "answer":"Phil Rudd"}
# {"question":"What is AC/DC's highest selling album in the United States?", "answer":"Back in Black"}
# {"question":"Approximately how many copies has Back in Black sold worldwide?", "answer":"Over 50 million"}
# {"question":"What famous AC/DC song begins with church bells?", "answer":"Hells Bells"}
# {"question":"In what year was the album The Razors Edge released?", "answer":"1990"}
# {"question":"Which AC/DC song features the lyrics 'I'm TNT, I'm dynamite'?", "answer":"T.N.T."}
# {"question":"What is the name of AC/DC's 2008 album?", "answer":"Black Ice"}
# {"question":"Which AC/DC song includes the line 'For those about to rock, we salute you'?", "answer":"For Those About to Rock (We Salute You)"}
# {"question":"When was the album Powerage released?", "answer":"1978"}
# {"question":"What is the name of AC/DC's 2020 album?", "answer":"Power Up"}
# {"question":"Who returned to AC/DC as vocalist for Power Up?", "answer":"Brian Johnson"}
# {"question":"Which AC/DC song is known for its iconic guitar intro and title track?", "answer":"Back in Black"}
# {"question":"Who directed many of AC/DC's early music videos?", "answer":"David Mallet"}
# {"question":"What AC/DC song was featured in the Iron Man 2 soundtrack?", "answer":"Shoot to Thrill"}
# {"question":"Which AC/DC live album was released in 1992?", "answer":"AC/DC Live"}
# {"question":"What is the name of AC/DC's concert film released in 2011?", "answer":"Live at River Plate"}
# {"question":"Which AC/DC song contains the lyric 'It's a long way to the top if you wanna rock 'n' roll'?", "answer":"It's a Long Way to the Top (If You Wanna Rock 'n' Roll)"}
# {"question":"Who played the bagpipes on 'It's a Long Way to the Top'?", "answer":"Bon Scott"}
# {"question":"Which member of AC/DC announced retirement due to hearing issues in 2016?", "answer":"Brian Johnson"}
# {"question":"Who temporarily replaced Brian Johnson on tour in 2016?", "answer":"Axl Rose"}
# {"question":"What AC/DC album features the track 'Thunderstruck'?", "answer":"The Razors Edge"}
# {"question":"In which country did AC/DC originate?", "answer":"Australia"}
# {"question":"Which AC/DC album cover features a cannon?", "answer":"For Those About to Rock"}
# {"question":"What is the name of the AC/DC track used as the title song for the 2012 film Iron Man 2?", "answer":"Highway to Hell"}
# {"question":"Which city is considered AC/DC's place of origin?", "answer":"Sydney"}
# {"question":"In what country was AC/DC formed?", "answer":"Australia"}
# {"question":"Who were the founding brothers of AC/DC?", "answer":"Angus Young and Malcolm Young"}
# {"question":"What style of rock is AC/DC most associated with?", "answer":"Hard rock"}
# {"question":"Which early AC/DC singer was replaced by Bon Scott?", "answer":"Dave Evans"}
# {"question":"What year did Brian Johnson join AC/DC?", "answer":"1980"}
# {"question":"Who was AC/DC's lead vocalist during the Highway to Hell era?", "answer":"Bon Scott"}
# {"question":"What is the title track of AC/DC's third studio album?", "answer":"Dirty Deeds Done Dirt Cheap"}
# {"question":"Which AC/DC album includes the song 'Touch Too Much'?", "answer":"Highway to Hell"}
# {"question":"Which AC/DC guitarist is famous for Angus Young's schoolboy persona?", "answer":"Angus Young"}
# {"question":"Who was AC/DC's primary rhythm guitarist until 2014?", "answer":"Malcolm Young"}
# {"question":"What year did AC/DC release the album Let There Be Rock?", "answer":"1977"}
# {"question":"Which album contains the track 'Whole Lotta Rosie'?", "answer":"Let There Be Rock"}
# {"question":"Who produced several of AC/DC's classic albums including Back in Black?", "answer":"Mutt Lange"}
# {"question":"Which AC/DC song famously opens with the tolling of a bell?", "answer":"Hells Bells"}
# {"question":"How many tracks are on the Back in Black album?", "answer":"10"}
# {"question":"Which AC/DC album was the first to feature Brian Johnson?", "answer":"Back in Black"}
# {"question":"Who wrote most of AC/DC's guitar riffs?", "answer":"Malcolm Young"}
# {"question":"What brand of guitar is Angus Young known for playing?", "answer":"Gibson SG"}
# {"question":"Which bassist joined AC/DC in 1977 and stayed for decades?", "answer":"Cliff Williams"}
# {"question":"What is the title of AC/DC's album released in 1981?", "answer":"For Those About to Rock"}
# {"question":"Which AC/DC song mentions dynamite in its lyrics?", "answer":"T.N.T."}
# {"question":"Which AC/DC record includes the track 'You Shook Me All Night Long'?", "answer":"Back in Black"}
# {"question":"When was AC/DC's album Flick of the Switch released?", "answer":"1983"}
# {"question":"What is the name of AC/DC's 1985 album?", "answer":"Fly on the Wall"}
# {"question":"Which drummer is most associated with AC/DC's classic era?", "answer":"Phil Rudd"}
# {"question":"Who played drums for AC/DC on the album Power Up?", "answer":"Phil Rudd"}
# {"question":"What is the title of the AC/DC album released in 1986 featuring 'Who Made Who'?", "answer":"Who Made Who"}
# {"question":"What soundtrack featured several AC/DC songs in 2010?", "answer":"Iron Man 2"}
# {"question":"Which album did AC/DC release in 1988?", "answer":"Blow Up Your Video"}
# {"question":"Which AC/DC song begins with the lyric 'Living easy, living free'?", "answer":"Highway to Hell"}
# {"question":"Which AC/DC tour promoted the album Black Ice?", "answer":"Black Ice World Tour"}
# {"question":"In which year was the album Black Ice released?", "answer":"2008"}
# {"question":"Which AC/DC album features the song 'Rock 'n' Roll Train'?", "answer":"Black Ice"}
# {"question":"What is the title of AC/DC's 1995 album?", "answer":"Ballbreaker"}
# {"question":"Which AC/DC album came out in 2000?", "answer":"Stiff Upper Lip"}
# {"question":"What year did Malcolm Young retire from AC/DC due to health issues?", "answer":"2014"}
# {"question":"Who filled in on rhythm guitar for Malcolm Young from 2014 onward?", "answer":"Stevie Young"}
# {"question":"What AC/DC compilation album was released in 1984?", "answer":"'74 Jailbreak"}
# {"question":"Which live album did AC/DC release in the early 1990s?", "answer":"AC/DC Live"}
# {"question":"What AC/DC track includes the phrase 'Thunder' shouted repeatedly?", "answer":"Thunderstruck"}
# {"question":"Which AC/DC record includes the song 'Moneytalks'?", "answer":"The Razors Edge"}
# {"question":"What is the name of the AC/DC track that references cannons?", "answer":"For Those About to Rock (We Salute You)"}
# {"question":"Which album includes the song 'Shoot to Thrill'?", "answer":"Back in Black"}
# {"question":"Who temporarily fronted AC/DC on tour when Brian Johnson stepped away in 2016?", "answer":"Axl Rose"}
# {"question":"What AC/DC album features the song 'Rock or Bust'?", "answer":"Rock or Bust"}
# {"question":"When was the album Rock or Bust released?", "answer":"2014"}
# {"question":"What is the title of AC/DC's 17th studio album?", "answer":"Power Up"}
# {"question":"Which AC/DC song mentions a cannon in its lyrics?", "answer":"For Those About to Rock (We Salute You)"}
# {"question":"Where was the music video for 'It's a Long Way to the Top' filmed?", "answer":"Melbourne"}
# {"question":"Which AC/DC song features bagpipes played by Bon Scott?", "answer":"It's a Long Way to the Top"}
# {"question":"Which AC/DC album cover shows a guitar with devil horns and a tail?", "answer":"Highway to Hell"}
# {"question":"What is the opening track of Highway to Hell?", "answer":"Highway to Hell"}
# {"question":"What instrument did Cliff Williams play in AC/DC?", "answer":"Bass guitar"}
# {"question":"Who was AC/DC's longest-serving drummer?", "answer":"Phil Rudd"}
# {"question":"Which AC/DC member is known for his energetic stage presence and duckwalk?", "answer":"Angus Young"}
# {"question":"How many siblings are in the Young family who were associated with AC/DC?", "answer":"Three"}
# {"question":"Which AC/DC co-founder passed away in 2017?", "answer":"Malcolm Young"}
# {"question":"What AC/DC release is considered one of the best-selling albums ever?", "answer":"Back in Black"}
# {"question":"Which AC/DC song has a title referring to electricity?", "answer":"High Voltage"}
# {"question":"Who handled bass duties before Cliff Williams joined AC/DC?", "answer":"Mark Evans"}
# {"question":"Which AC/DC album title includes the word 'Voltage'?", "answer":"High Voltage"}
# {"question":"Who is known for writing many of AC/DC's lyrics during the Bon Scott era?", "answer":"Bon Scott"}
# {"question":"Which AC/DC album contains the song 'Problem Child'?", "answer":"Dirty Deeds Done Dirt Cheap"}
# {"question":"What is the name of the AC/DC track with the lyric 'I'm back in black'?", "answer":"Back in Black"}
# {"question":"Who managed AC/DC during their early international success?", "answer":"Michael Browning"}
# {"question":"Which AC/DC album was promoted with the Blow Up Your Video World Tour?", "answer":"Blow Up Your Video"}
# {"question":"Which AC/DC song includes the line 'She told me to come but I was already there'?", "answer":"You Shook Me All Night Long"}
# {"question":"What is the name of the AC/DC track used frequently at sporting events?", "answer":"Thunderstruck"}
# {"question":"Which AC/DC album features the track 'Heatseeker'?", "answer":"Blow Up Your Video"}
# {"question":"What AC/DC release contains the song 'Rock 'n' Roll Damnation'?", "answer":"Powerage"}
# {"question":"Who played drums on AC/DC's earliest recordings?", "answer":"Colin Burgess"}
# {"question":"What was the first AC/DC album released after Bon Scott's death?", "answer":"Back in Black"}
# {"question":"Which AC/DC song features cannons firing during live performances?", "answer":"For Those About to Rock (We Salute You)"}
# {"question":"What year was the AC/DC album Power Up issued?", "answer":"2020"}
# {"question":"Which AC/DC album includes the track 'Shot in the Dark'?", "answer":"Power Up"}
# {"question":"What is the opening song on the album Back in Black?", "answer":"Hells Bells"}
# {"question":"When was Cold Chisel formed?", "answer":"1973"}
# {"question":"In which country did Cold Chisel originate?", "answer":"Australia"}
# {"question":"Who was the lead singer of Cold Chisel?", "answer":"Jimmy Barnes"}
# {"question":"Who played piano and keyboards in Cold Chisel?", "answer":"Don Walker"}
# {"question":"Which city was home base for Cold Chisel in their early years?", "answer":"Adelaide"}
# {"question":"What genre is Cold Chisel primarily known for?", "answer":"Pub rock"}
# {"question":"What is Cold Chisel's best-known live album?", "answer":"Swingshift"}
# {"question":"Which Cold Chisel song includes the lyric 'Standing on the outside looking in'?", "answer":"Standing on the Outside"}
# {"question":"Who was the guitarist for Cold Chisel?", "answer":"Ian Moss"}
# {"question":"What instrument did Steve Prestwich play in Cold Chisel?", "answer":"Drums"}
# {"question":"When did Cold Chisel release their debut album?", "answer":"1978"}
# {"question":"What is the name of Cold Chisel's debut album?", "answer":"Cold Chisel"}
# {"question":"Which album contains the song 'Khe Sanh'?", "answer":"Cold Chisel"}
# {"question":"What Cold Chisel track references Vietnam War veterans?", "answer":"Khe Sanh"}
# {"question":"Which album includes the song 'Flame Trees'?", "answer":"Twentieth Century"}
# {"question":"When did Cold Chisel break up for the first time?", "answer":"1983"}
# {"question":"Who wrote many of Cold Chisel's lyrics?", "answer":"Don Walker"}
# {"question":"Which Cold Chisel album features 'Choirgirl'?", "answer":"East"}
# {"question":"When was the album East released?", "answer":"1980"}
# {"question":"Who was the bassist for Cold Chisel?", "answer":"Phil Small"}
# {"question":"Which Cold Chisel song includes the line 'Here comes the working man'?", "answer":"Working Class Man"}
# {"question":"What year did Cold Chisel reunite for The Last Stand tour film?", "answer":"1983"}
# {"question":"Which Cold Chisel album was released in 2012?", "answer":"No Plans"}
# {"question":"What is the title of their 2015 album?", "answer":"The Perfect Crime"}
# {"question":"Which member of Cold Chisel passed away in 2011?", "answer":"Steve Prestwich"}
# {"question":"Who founded the Foo Fighters?", "answer":"Dave Grohl"}
# {"question":"When was Foo Fighters formed?", "answer":"1994"}
# {"question":"Which city is the origin of Foo Fighters?", "answer":"Seattle"}
# {"question":"What instrument does Dave Grohl primarily play in Foo Fighters?", "answer":"Guitar and vocals"}
# {"question":"Which former Nirvana member created Foo Fighters?", "answer":"Dave Grohl"}
# {"question":"What was Foo Fighters' debut album called?", "answer":"Foo Fighters"}
# {"question":"In what year was the album The Colour and the Shape released?", "answer":"1997"}
# {"question":"Which Foo Fighters album features the song 'Everlong'?", "answer":"The Colour and the Shape"}
# {"question":"What is the name of Foo Fighters' 2011 album?", "answer":"Wasting Light"}
# {"question":"Who was Foo Fighters' longtime drummer until 2022?", "answer":"Taylor Hawkins"}
# {"question":"Which Foo Fighters track starts with the lyric 'Hello, I've waited here for you'?", "answer":"Everlong"}
# {"question":"What is the title of Foo Fighters' 2005 double album?", "answer":"In Your Honor"}
# {"question":"Which album includes the song 'Best of You'?", "answer":"In Your Honor"}
# {"question":"Who plays bass for Foo Fighters?", "answer":"Nate Mendel"}
# {"question":"What is the title of the Foo Fighters' 1999 album?", "answer":"There Is Nothing Left to Lose"}
# {"question":"Which Foo Fighters song features the lyric 'Is someone getting the best of you'?", "answer":"Best of You"}
# {"question":"Which Foo Fighters album contains 'Learn to Fly'?", "answer":"There Is Nothing Left to Lose"}
# {"question":"When was the album Concrete and Gold released?", "answer":"2017"}
# {"question":"What is the name of the Foo Fighters' 2021 album?", "answer":"Medicine at Midnight"}
# {"question":"Which Foo Fighters member was originally in the band Sunny Day Real Estate?", "answer":"Nate Mendel"}
# {"question":"What is the name of the Foo Fighters' 2002 album?", "answer":"One by One"}
# {"question":"Who directed the documentary Sonic Highways?", "answer":"Dave Grohl"}
# {"question":"Which Foo Fighters album accompanied the Sonic Highways series?", "answer":"Sonic Highways"}
# {"question":"Which original Foo Fighters member died in 2022?", "answer":"Taylor Hawkins"}
# {"question":"What is the title of Foo Fighters' debut single?", "answer":"This Is a Call"}
#

"""Train the model"""

def train_model(model_name, data_files, output_dir, truncation, batch_size, num_epochs, learning_rate):

  # Load the training data

  #model_name = "google/flan-t5-small"
  tokenizer = AutoTokenizer.from_pretrained(model_name)
  dataset = load_dataset("json", data_files=data_files)

  def preprocess(example):
    inputs = tokenizer(example["question"], max_length=128, truncation=truncation, padding="max_length")
    targets = tokenizer(example["answer"], max_length=128, truncation=truncation, padding="max_length")

    #inputs["labels"] = targets["input_ids"]

    labels = targets["input_ids"]
    labels = [
        label if label != tokenizer.pad_token_id else -100
        for label in labels
    ]
    inputs["labels"] = labels

    return inputs

  tokenized_dataset = dataset.map(preprocess)
  tokenized_dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

  # Prepare the model for training

  model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

  peft_config = LoraConfig(
    task_type=TaskType.SEQ_2_SEQ_LM,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1
  )

  model = get_peft_model(model, peft_config)

  # Configure and run the trainer

  training_args = Seq2SeqTrainingArguments(
    output_dir=output_dir,
    per_device_train_batch_size=batch_size,
    num_train_epochs=num_epochs,
    logging_steps=1,
    push_to_hub=False,
    learning_rate=learning_rate,
    eval_strategy="epoch",
  )

  trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset['train'],
    eval_dataset=tokenized_dataset['train'].select(range(len(tokenized_dataset['train']))),
  )

  trainer.train()

  # Save the fine-tuned model and tokenizer

  model.save_pretrained(output_dir)
  tokenizer.save_pretrained(output_dir)

"""Test the trained model"""

def test():

  qa_pipeline = pipeline("text2text-generation", model=output_dir, tokenizer=output_dir)

  questions = [
    "When was AC/DC formed?",
    "Where was AC/DC formed?",
    "List the members of Cold Chisel.",
    "List the members of AC/DC.",
  ]

  for question in questions:
    answer = qa_pipeline(question)
    print(f"{question} Answer: {answer[0]['generated_text']}")

"""Parameters"""

model_name = "google/flan-t5-small"
data_files = "acdc_qa.json"
output_dir = "./acdc-finetuned-model"
truncation = False
batch_size = 4
num_epochs = 100
learning_rate = 1e-3

"""Train the model"""

train_model(model_name, data_files, output_dir, truncation, batch_size, num_epochs, learning_rate)

"""Test the trained model"""

test()