mirror of
https://github.com/wassname/Open-Assistant.git
synced 2026-07-19 11:20:04 +08:00
- standardize file naming to be a little more pythonic. - one folder per notebook plus a README for easy GH viewing. - adjust colab badges for new structure.
2.6 MiB
2.6 MiB
In [ ]:
# uncomment below to install required python packages
#!pip install detoxifyIn [14]:
from detoxify import Detoxify
import gc
import torch
import requests
import random
import seaborn as sns
import pandas as pd
import matplotlib.pyplot as plt
import timeIn [15]:
detoxify_models = ["original", "original-small", "multilingual", "unbiased", "unbiased-small"]
not_obvious_toxic = [
"I'm not a racist, but I don't like afroamerican people",
"I love my girlfriend, but I don't like her black friends",
"My girfriend is the best and most loving person in the world as long as she doesn't leave the kitchen",
"My girlfriend is 12 years old but I love her, after all age is just a number",
"You are like a star to me, you should stay milion of light years away",
]
not_obvious_non_toxic = [
"I fucking love you motherfucker",
"I love my girlfriend so much I would kill for her",
"Wow you have an important event soon? Break a leg!",
"I'm 35 and my girlfriend is 12. Can we do quests together with such a huge level gap?",
"I have reat this really awesome book about barbarians slaughtering villages and raping women",
]
obvious_toxic = [
"I absolutly hate you man",
"This question is pathetic, you are an idiot",
"I hate this minority and I think they should be exterminated",
"Women should only work in the kitchen and keep the house clean",
"Men should always be muscular and never cry ",
]
obvious_non_toxic = [
"I love my girlfriend, she means a lot to me",
"You are a very good friend and you mean a lot to me",
"What are we having for dinner tonight?",
"People should read more books as they are very educating",
"We just returned from our trip to Spain",
]
word_site = "https://www.mit.edu/~ecprice/wordlist.10000"
response = requests.get(word_site)
WORDS = [word.decode("utf-8") for word in response.content.splitlines()]
DEVICE = "cuda"In [16]:
def random_sentence(sentence_length):
return " ".join([WORDS[random.randint(0, len(WORDS) - 1)] for i in range(sentence_length)])In [ ]:
for model in detoxify_models:
print(f"Loading {model} model")
Detoxify(model)
gc.collect()
print(f"Loaded {model} model")In [18]:
def check_model(model_name):
torch.cuda.empty_cache()
initial_memory = torch.cuda.memory_allocated()
model = Detoxify(model_name, device=DEVICE)
model_memory = (torch.cuda.memory_allocated() - initial_memory) / (1024 * 1024)
max_sentence_length = 4000
max_batch_size = 128
sentence_step = 500
batch_step = 32
memory_heatmap = pd.DataFrame(
columns=[i for i in range(sentence_step, max_sentence_length + 1, sentence_step)],
index=[i for i in range(batch_step, max_batch_size + 1, batch_step)],
)
execution_time_heatmap = pd.DataFrame(
columns=[i for i in range(sentence_step, max_sentence_length + 1, sentence_step)],
index=[i for i in range(batch_step, max_batch_size + 1, batch_step)],
)
for word_size in range(sentence_step, max_sentence_length + 1, sentence_step):
for batch_size in range(batch_step, max_batch_size + 1, batch_step):
start_time = time.time()
inputs = [random_sentence(word_size) for i in range(batch_size)]
_ = model.predict(inputs)
memory_heatmap.loc[batch_size, word_size] = (torch.cuda.max_memory_allocated() - initial_memory) / (
1024 * 1024
)
execution_time_heatmap.loc[batch_size, word_size] = time.time() - start_time
del inputs, _
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
plt.figure(figsize=(20, 20))
plt.suptitle(f'Detoxify model "{model_name}" base memory usage = {model_memory:.2f} MB', fontsize=36)
plt.subplot(2, 2, 1)
sns.heatmap(memory_heatmap.astype(float), annot=True, fmt=".0f", cmap="Blues")
plt.title(f"{model_name} model inference memory usage (MB)")
plt.xlabel("Sentence length")
plt.ylabel("Batch size")
plt.subplot(2, 2, 2)
sns.heatmap(execution_time_heatmap.astype(float), annot=True, fmt=".2f", cmap="Blues")
plt.title(f"{model_name} model inference execution time (seconds)")
plt.xlabel("Sentence length")
plt.ylabel("Batch size")
max_sentence_length = 4000
max_batch_size = 16
sentence_step = 500
batch_step = 4
memory_heatmap = pd.DataFrame(
columns=[i for i in range(sentence_step, max_sentence_length + 1, sentence_step)],
index=[i for i in range(batch_step, max_batch_size + 1, batch_step)],
)
execution_time_heatmap = pd.DataFrame(
columns=[i for i in range(sentence_step, max_sentence_length + 1, sentence_step)],
index=[i for i in range(batch_step, max_batch_size + 1, batch_step)],
)
optimizer = torch.optim.Adam(model.model.parameters(), lr=0.0001)
for word_size in range(sentence_step, max_sentence_length + 1, sentence_step):
for batch_size in range(batch_step, max_batch_size + 1, batch_step):
model.model.train()
start_time = time.time()
inputs = [random_sentence(word_size) for i in range(batch_size)]
outputs = model.model(
**model.tokenizer(inputs, return_tensors="pt", padding=True, truncation=True).to(DEVICE)
)[0]
outputs = torch.sigmoid(outputs)
random_outputs = torch.rand(outputs.shape).to(DEVICE)
loss = torch.nn.functional.binary_cross_entropy(outputs, random_outputs)
loss.backward()
optimizer.step()
memory_heatmap.loc[batch_size, word_size] = (torch.cuda.max_memory_allocated() - initial_memory) / (
1024 * 1024
)
execution_time_heatmap.loc[batch_size, word_size] = time.time() - start_time
del inputs, outputs, random_outputs, loss
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
plt.subplot(2, 2, 3)
sns.heatmap(memory_heatmap.astype(float), annot=True, fmt=".0f", cmap="Blues")
plt.title(f"{model_name} model training memory usage (MB)")
plt.xlabel("Sentence length")
plt.ylabel("Batch size")
plt.subplot(2, 2, 4)
sns.heatmap(execution_time_heatmap.astype(float), annot=True, fmt=".2f", cmap="Blues")
plt.title(f"{model_name} model training execution time (seconds)")
plt.xlabel("Sentence length")
plt.ylabel("Batch size")
for m in detoxify_models:
check_model(m)In [19]:
def check_outputs(model_name):
model = Detoxify(model_name, device=DEVICE)
should_be_toxic = pd.DataFrame(model.predict(not_obvious_toxic))
should_not_be_toxic = pd.DataFrame(model.predict(not_obvious_non_toxic))
must_be_toxic = pd.DataFrame(model.predict(obvious_toxic))
must_not_be_toxic = pd.DataFrame(model.predict(obvious_non_toxic))
nl = "\n" # f strings don't support new lines
plt.figure(figsize=(15, 15))
plt.suptitle(f'Detoxify model "{model_name}" outputs', fontsize=30)
plt.subplot(2, 2, 1)
sns.heatmap(should_be_toxic, annot=True, fmt=".2f", cmap="Blues")
plt.title(f'not obvious toxic {nl} { "".join([f"{i}: {s} {nl}" for i, s in enumerate(not_obvious_toxic)])}')
plt.subplot(2, 2, 2)
sns.heatmap(should_not_be_toxic, annot=True, fmt=".2f", cmap="Blues")
plt.title(f'not obvious not toxic {nl} { "".join([f"{i}: {s} {nl}" for i, s in enumerate(not_obvious_non_toxic)])}')
plt.subplot(2, 2, 3)
sns.heatmap(must_be_toxic, annot=True, fmt=".2f", cmap="Blues")
plt.title(f'obvious toxic {nl} { "".join([f"{i}: {s} {nl}" for i, s in enumerate(obvious_toxic)])}')
plt.subplot(2, 2, 4)
sns.heatmap(must_not_be_toxic, annot=True, fmt=".2f", cmap="Blues")
plt.title(f'obvious not toxic {nl} { "".join([f"{i}: {s} {nl}" for i, s in enumerate(obvious_non_toxic)])}')
plt.tight_layout()
for m in detoxify_models:
check_outputs(m)