Generative AI - Building a Chatbot with Transformer Models

Building a Chatbot with Transformer Models

In this project, we will create a chatbot using the Transformer model and the Hugging Face Transformers library. We will use a pre-trained model and fine-tune it on a bespoke dataset to ensure that the chatbot responds correctly to user input.

Step 1: Setup and Import Necessary Libraries

First, we must install and import the essential libraries.

# Install necessary libraries
!pip install transformers torch datasets

# Import libraries
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from datasets import load_dataset, Dataset

We use pip to install the Hugging Face Transformers, PyTorch, and datasets libraries.

We import the relevant modules from transformers, torch, and datasets.

Step 2: Load and Prepare the Dataset

We will train our chatbot using the Cornell Movie Dialogues dataset.

# Load the dataset
dataset = load_dataset('cornell_movie_dialogue')

# Extract dialogues and prepare the data
def preprocess_data(dialogues):
    inputs = []
    outputs = []
    for i in range(len(dialogues) - 1):
        inputs.append(dialogues[i]['text'])
        outputs.append(dialogues[i + 1]['text'])
    return inputs, outputs

dialogues = dataset['train']['dialogue']
inputs, outputs = preprocess_data(dialogues)

# Create a Dataset object
data = {'input_text': inputs, 'output_text': outputs}
dataset = Dataset.from_dict(data)

The Cornell Movie Dialogues dataset is loaded using the load_dataset function.

We define the preprocess_data function to extract conversations and create input-output pairs.

We generate a dictionary of input and output text and convert it to a Dataset object.

Step 3: Tokenize the Data

We'll use a pre-trained tokenizer to tokenize the text data.

# Load the pre-trained tokenizer
tokenizer = AutoTokenizer.from_pretrained('gpt2')

# Tokenize the data
def tokenize_function(example):
    return tokenizer(example['input_text'], padding='max_length', truncation=True)

tokenized_dataset = dataset.map(tokenize_function, batched=True)

We use the pre-trained tokenizer from the 'gpt2' model.

We create a function called tokenize_function to tokenize the input text.

We use the tokenization function on the dataset.

Step 4: Initialize the Transformer Model

We set up the transformer model for causal language modeling.

# Load the pre-trained model
model = AutoModelForCausalLM.from_pretrained('gpt2')

We use the pre-trained GPT-2 model for causal language modeling.

Step 5: Define Training Arguments and Trainer

We describe the reasons for teaching and setting up the Trainer.

# Define training arguments
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=10,
)

# Initialize the Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    eval_dataset=tokenized_dataset,
)

The training inputs include the output path, the number of epochs, the batch sizes, the warmup steps, the rate at which the weights decrease, and the setup for logging.

The model, training inputs, and tokenized dataset are used to set up the Trainer.

Step 6: Train and Evaluate the Model

We will use the Trainer to teach and test the model.

# Train the model
trainer.train()

# Evaluate the model
results = trainer.evaluate()
print(results)

We use the training data to fine-tune the transformer model using the training method of the Trainer.

We use the evaluation method to assess the model and report the results.

Step 7: Interact with the Chatbot

We'll create a function to communicate with the trained chatbot.

# Function to generate responses from the chatbot
def generate_response(model, tokenizer, input_text):
    input_ids = tokenizer.encode(input_text, return_tensors='pt')
    output = model.generate(input_ids, max_length=50, num_return_sequences=1)
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response

# Interact with the chatbot
while True:
    user_input = input("You: ")
    if user_input.lower() == 'exit':
        break
    response = generate_response(model, tokenizer, user_input)
    print(f"Chatbot: {response}")

We define the code generate_response so that the robot can send replies.

The method encodes the text that is passed in, uses the model to make a reaction, and then decodes the result.

We make a loop for the user to connect with the robot, which lets them type text and get answers.