[干货] 一文介绍如何训练GPT2，让自己的数据会说话

Company Profile:

Company Name: Dummy-Gpt2-Datatec-Studio Inc

At Dummy-Gpt2-Datatec-Studio Inc, we are a cutting-edge technology company committed to transforming visionary ideas into tangible realities. Our mission is to pioneer innovation that makes a meaningful impact on the world, from extending human longevity to addressing critical environmental challenges and revolutionizing the way we travel.

Portfolio:

Life-Enhancing Pharmaceuticals: Our portfolio includes groundbreaking pharmaceutical research aimed at developing vaccines and treatments to prevent diseases, potentially adding an additional 20 years to the human lifespan.
Environmental Sustainability: We are leaders in developing innovative solutions to combat climate change and ensure a sustainable future for the planet, encompassing green energy, waste reduction, and carbon footprint mitigation.
Advanced Transportation: In our transportation division, we engineer state-of-the-art terrestrial transportation solutions, promising ultra-fast and efficient modes of travel that will reshape the way we move.
Products:

Green Energy Devices: Our product line features advanced devices that empower individuals and businesses to generate clean and renewable energy, rendering traditional power plants and fossil fuels obsolete.
Space Exploration Services: Our space exploration services offer accessible and safe travel options to explore colonies on celestial bodies such as solar planets and planetary moons, opening new frontiers for exploration.
Quantum Computing Solutions: We provide cutting-edge quantum computing solutions, exponentially boosting computational power and network speeds, thus ushering in a new era of technological advancement.
Holographic Communication: Our holographic communication technology replaces traditional televisions and phones with immersive, interactive experiences. Integrating the senses of touch and smell, we bring the virtual world closer to reality.
Services:

Genetic Engineering: Our genetic engineering services explore revolutionary concepts like Wi-Fi-emitting plants, potentially transforming the way we connect and communicate in a lush and connected world.
Futuristic Planning Consultation: We specialize in assisting individuals and organizations in planning for the distant future. Leveraging transformative technologies, we help shape and anticipate possible outcomes to achieve long-term goals.
Contact:

For inquiries and collaborations, please reach out to us at:

Email: contact@dummygpt2datatecstudio.com
Phone: +1-123-456-7890
Address: 123 Innovator's Lane, Futuroville, Earth
Community Engagement:

At Dummy-Gpt2-Datatec-Studio Inc, we take our commitment to community involvement seriously. We actively support STEM education programs, empowering future generations to be at the forefront of scientific and technological advancements that will shape our collective future.

This updated description provides a more concrete and detailed overview of Dummy-Gpt2-Datatec-Studio Inc, reflecting its visionary pursuits and the potential impact on various aspects of life and technology.

os.environ["HF_HOME_TOKEN"] = "Please_replace_it_with_your_hf_access_token"

[
  {
    "dialog_id": 1,
    "dialog": [
      {
        "id": 0,
        "sender": "user",
        "text": "Tell me about your company."
      },
      {
        "id": 1,
        "sender": "company",
        "text": "We are Dummy-Gpt2-Datatec-Studio Inc, a cutting-edge technology company dedicated to pioneering innovation."
      }
    ],
    "eval_score": 4,
    "profile_match": 5
  },
......
]

# Fine tuning gpt2_medium model and use own data like company profile
#
# See also medium.com blog
# "GPT-2 Fine-Tuning Guide: Building a Chatbot for Your Company Profile"
# https://medium.com/@datatec.studio
#

# Mount google driver
from google.colab import drive
drive.mount('/content/drive')

# Change to google driver folder which contains datasets
# This folder will also be used to save model
print("Please upload the github dataset to your google driver folder GPT2_Lab_DTS")
print("github repository: https://github.com/datatecyl/gpt2_lab/tree/master/google_driver")
%cd /content/drive/MyDrive/GPT2_Lab_DTS

# Install related python package. The requirements.txt is from google driver.
!pip install -r requirements.txt

# Import packages
import os
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch
from torch.optim import Adam
from torch.utils.data import DataLoader
import tqdm

# Define path and environment
result_dir = '/content/drive/MyDrive/GPT2_Lab_DTS/results'
data_file_path = '/content/drive/MyDrive/GPT2_Lab_DTS/data/my_company_info.json'
os.environ["HF_HOME"] = "/content/huggingface"  # Replace with your desired directory
print("Please replace it with your hf access token:")
os.environ["HF_HOME_TOKEN"] = "Please_replace_it_with_your_hf_access_token"

model_name = "gpt2-medium"
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(device)

# Write a python file to google driver
# Sample of json datasets
# You can also directly upload this code to your google driver
# The code write here in this way is for better understanding of whole project
%%writefile chat_data.py

from torch.utils.data import Dataset
import json

class ChatData(Dataset):
    def __init__(self, path: str, tokenizer):
        self.data = json.load(open(path, "r"))

        self.X = []
        for i in self.data:
            for j in i['dialog']:
                self.X.append(j['text'])

        for idx, i in enumerate(self.X):
            try:
                self.X[idx] = "<startofstring> " + i + " <bot>: " + self.X[idx + 1] + " <endofstring>"
            except:
                break

        for i in self.data:
            for j in i['dialog']:
                self.X.append(j['text'])

        total_samples = len(self.X)  # Calculate the total number of samples
        print("total_samples", total_samples)
        # define samples amount
        self.X = self.X[:500]
        print("Here is the self.X[0] i wanna check:")
        print(self.X[0])

        self.X_encoded = tokenizer(self.X, return_tensors="pt", max_length=30, padding="max_length", truncation=True)
        self.input_ids = self.X_encoded['input_ids']
        self.attention_mask = self.X_encoded['attention_mask']

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.input_ids[idx], self.attention_mask[idx]


# Download model, save model and tokernize to harddisk
## prepare tokenizer
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
tokenizer.add_special_tokens({"pad_token": "<pad>",
                            "bos_token": "<startofstring>",
                            "eos_token": "<endofstring>"})
tokenizer.add_tokens(["<bot>:"])

## prepare model
### Specify the desired embedding size (must be a multiple of 8)
desired_embedding_size = 50264  # Change this to the desired size
model = GPT2LMHeadModel.from_pretrained(model_name).to(device)
### Resize the embedding layer to the desired size
model.resize_token_embeddings(len(tokenizer), desired_embedding_size)
model = model.to(device)

## save tokenizer and model to harddisk
tokenizer.save_pretrained(result_dir)
model.save_pretrained(result_dir)

## load model and tokenizer from harddisk
### Load the GPT-2 tokenizer
tokenizer = GPT2Tokenizer.from_pretrained(result_dir)


### Load the GPT-2 model from the local folder
model = GPT2LMHeadModel.from_pretrained(result_dir)
model.to(device)

# Define infer and train function
def infer(inp):
  inp = "<startofstring> " + inp + " <bot>: "
  inp = tokenizer(inp, return_tensors="pt")
  X = inp["input_ids"].to(device)  # Use .to(device) method to move the tensor to the specified device
  a = inp["attention_mask"].to(device)  # Use .to(device) method here as well

  output = model.generate(X, attention_mask=a, max_length=100, num_return_sequences=1)

  output = tokenizer.decode(output[0])

  return output

def train(chatData, model, optim):

  epochs = 12

  for _ in tqdm.tqdm(range(epochs)):  # Use range() to iterate through epochs
      for X, a in chatData:
          print(X)
          X = X.to(device)
          a = a.to(device)
          optim.zero_grad()
          loss = model(input_ids=X, attention_mask=a, labels=X).loss
          loss.backward()
          optim.step()

  # Save the model's state dictionary after training is complete
  torch.save(model.state_dict(), "model_state.pt")
  print(infer("How do you see the integration of holographic technology in education?"))

# Load ChatData, train model and optimizer
from chat_data import ChatData

chatData = ChatData(data_file_path, tokenizer)
chatData = DataLoader(chatData, batch_size=64)

model.train()

optim = Adam(model.parameters())

# train 10 times

epochs = 10  # You can adjust the number of epochs as needed
for epoch in range(epochs):
    print("Round: ", epoch)
    train(chatData, model, optim)

# Show input textfield to interaction with model

inp = ""
while True:
  inp = input("Enter your input (press Enter when done): " + " " * 20)
  print(infer(inp))