mirror of
https://github.com/twitter/the-algorithm.git
synced 2024-09-30 19:08:40 +02:00
325 lines
9.2 KiB
Python
325 lines
9.2 KiB
Python
import tensorflow as tf
|
|
|
|
physical_devices = tf.config.list_physical_devices("GPU")
|
|
for device in physical_devices:
|
|
tf.config.experimental.set_memory_growth(device, True)
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import utils
|
|
from twitter.cuad.representation.models.optimization import create_optimizer
|
|
from twitter.cuad.representation.models.text_encoder import TextEncoder
|
|
from twitter.hmli.nimbus.modeling.feature_encoder import FeatureEncoder
|
|
from twitter.hmli.nimbus.modeling.feature_loader import BigQueryFeatureLoader
|
|
from twitter.hmli.nimbus.modeling.model_config import (
|
|
EncodingType,
|
|
Feature,
|
|
FeatureType,
|
|
LogType,
|
|
Model,
|
|
)
|
|
|
|
cat_names = [...]
|
|
|
|
category_features = [
|
|
Feature(name=cat_name, ftype=FeatureType.CONTINUOUS) for cat_name in cat_names
|
|
]
|
|
features = [
|
|
Feature(
|
|
name="tweet_text_with_media_annotations",
|
|
ftype=FeatureType.STRING,
|
|
encoding=EncodingType.BERT,
|
|
),
|
|
Feature(name="precision_nsfw", ftype=FeatureType.CONTINUOUS),
|
|
Feature(name="has_media", ftype=FeatureType.BINARY),
|
|
Feature(name="num_media", ftype=FeatureType.DISCRETE),
|
|
] + category_features
|
|
|
|
ptos_prototype = Model(
|
|
name="ptos_prototype",
|
|
export_path="...",
|
|
features=features,
|
|
)
|
|
print(ptos_prototype)
|
|
|
|
cq_loader = BigQueryFeatureLoader(gcp_project=COMPUTE_PROJECT)
|
|
labels = [
|
|
"has_non_punitive_action",
|
|
"has_punitive_action",
|
|
"has_punitive_action_contains_self_harm",
|
|
"has_punitive_action_encourage_self_harm",
|
|
"has_punitive_action_episodic",
|
|
"has_punitive_action_episodic_hateful_conduct",
|
|
"has_punitive_action_other_abuse_policy",
|
|
"has_punitive_action_without_self_harm",
|
|
]
|
|
|
|
train_query = f"""
|
|
SELECT
|
|
{{feature_names}},
|
|
{",".join(labels)},
|
|
...
|
|
"""
|
|
val_query = f"""
|
|
SELECT
|
|
{{feature_names}},
|
|
{",".join(labels)},
|
|
...
|
|
"""
|
|
|
|
print(train_query)
|
|
train = cq_loader.load_features(ptos_prototype, "", "", custom_query=train_query)
|
|
val = cq_loader.load_features(ptos_prototype, "", "", custom_query=val_query)
|
|
print(train.describe(model=ptos_prototype))
|
|
|
|
params = {
|
|
"max_seq_lengths": 128,
|
|
"batch_size": 196,
|
|
"lr": 1e-5,
|
|
"optimizer_type": "adamw",
|
|
"warmup_steps": 0,
|
|
"cls_dropout_rate": 0.1,
|
|
"epochs": 30,
|
|
"steps_per_epoch": 5000,
|
|
"model_type": "twitter_multilingual_bert_base_cased_mlm",
|
|
"mixed_precision": True,
|
|
}
|
|
params
|
|
|
|
|
|
def parse_labeled_data(row_dict):
|
|
label = [row_dict.pop(l) for l in labels]
|
|
return row_dict, label
|
|
|
|
|
|
mirrored_strategy = tf.distribute.MirroredStrategy()
|
|
BATCH_SIZE = params["batch_size"] * mirrored_strategy.num_replicas_in_sync
|
|
|
|
train_ds = (
|
|
train.to_tf_dataset()
|
|
.map(parse_labeled_data)
|
|
.shuffle(BATCH_SIZE * 100)
|
|
.batch(BATCH_SIZE)
|
|
.repeat()
|
|
)
|
|
val_ds = val.to_tf_dataset().map(parse_labeled_data).batch(BATCH_SIZE)
|
|
|
|
for record in train_ds:
|
|
tf.print(record)
|
|
break
|
|
|
|
|
|
def get_positive_weights():
|
|
"""Computes positive weights used for class imbalance from training data."""
|
|
label_weights_df = utils.get_label_weights(
|
|
"tos-data-media-full",
|
|
project_id="twttr-abusive-interact-prod",
|
|
dataset_id="tos_policy",
|
|
)
|
|
pos_weight_tensor = tf.cast(
|
|
label_weights_df.sort_values(by="label").positive_class_weight, dtype=tf.float32
|
|
)
|
|
return pos_weight_tensor
|
|
|
|
|
|
pos_weight_tensor = get_positive_weights()
|
|
print(pos_weight_tensor)
|
|
|
|
|
|
class TextEncoderPooledOutput(TextEncoder):
|
|
def call(self, x):
|
|
return super().call([x])["pooled_output"]
|
|
|
|
def get_config(self):
|
|
return super().get_config()
|
|
|
|
|
|
with mirrored_strategy.scope():
|
|
text_encoder_pooled_output = TextEncoderPooledOutput(
|
|
params["max_seq_lengths"], model_type=params["model_type"], trainable=True
|
|
)
|
|
|
|
fe = FeatureEncoder(train)
|
|
inputs, preprocessing_head = fe.build_model_head(
|
|
model=ptos_prototype, text_encoder=text_encoder_pooled_output
|
|
)
|
|
|
|
cls_dropout = tf.keras.layers.Dropout(
|
|
params["cls_dropout_rate"], name="cls_dropout"
|
|
)
|
|
outputs = cls_dropout(preprocessing_head)
|
|
outputs = tf.keras.layers.Dense(8, name="output", dtype="float32")(outputs)
|
|
|
|
model = tf.keras.Model(inputs=inputs, outputs=outputs)
|
|
pr_auc = tf.keras.metrics.AUC(
|
|
curve="PR", num_thresholds=1000, multi_label=True, from_logits=True
|
|
)
|
|
|
|
custom_loss = lambda y_true, y_pred: utils.multilabel_weighted_loss(
|
|
y_true, y_pred, weights=pos_weight_tensor
|
|
)
|
|
optimizer = create_optimizer(
|
|
init_lr=params["lr"],
|
|
num_train_steps=(params["epochs"] * params["steps_per_epoch"]),
|
|
num_warmup_steps=params["warmup_steps"],
|
|
optimizer_type=params["optimizer_type"],
|
|
)
|
|
if params.get("mixed_precision"):
|
|
optimizer = tf.train.experimental.enable_mixed_precision_graph_rewrite(
|
|
optimizer
|
|
)
|
|
|
|
model.compile(optimizer=optimizer, loss=custom_loss, metrics=[pr_auc])
|
|
|
|
model.weights
|
|
model.summary()
|
|
pr_auc.name
|
|
|
|
import getpass
|
|
|
|
import wandb
|
|
from wandb.keras import WandbCallback
|
|
|
|
try:
|
|
wandb_key = ...
|
|
wandb.login(...)
|
|
run = wandb.init(
|
|
project="ptos_with_media",
|
|
group="new-split-trains",
|
|
notes="tweet text with only (num_media, precision_nsfw). on full train set, new split.",
|
|
entity="absv",
|
|
config=params,
|
|
name="tweet-text-w-nsfw-1.1",
|
|
sync_tensorboard=True,
|
|
)
|
|
except FileNotFoundError:
|
|
print("Wandb key not found")
|
|
run = wandb.init(mode="disabled")
|
|
import datetime
|
|
import os
|
|
|
|
start_train_time = datetime.datetime.now()
|
|
print(start_train_time.strftime("%m-%d-%Y (%H:%M:%S)"))
|
|
checkpoint_path = os.path.join("...")
|
|
print("Saving model checkpoints here: ", checkpoint_path)
|
|
|
|
cp_callback = tf.keras.callbacks.ModelCheckpoint(
|
|
filepath=os.path.join(checkpoint_path, "model.{epoch:04d}.tf"),
|
|
verbose=1,
|
|
monitor=f"val_{pr_auc.name}",
|
|
mode="max",
|
|
save_freq="epoch",
|
|
save_best_only=True,
|
|
)
|
|
|
|
early_stopping_callback = tf.keras.callbacks.EarlyStopping(
|
|
patience=7, monitor=f"val_{pr_auc.name}", mode="max"
|
|
)
|
|
|
|
model.fit(
|
|
train_ds,
|
|
epochs=params["epochs"],
|
|
validation_data=val_ds,
|
|
callbacks=[cp_callback, early_stopping_callback],
|
|
steps_per_epoch=params["steps_per_epoch"],
|
|
verbose=2,
|
|
)
|
|
|
|
import tensorflow_hub as hub
|
|
|
|
gs_model_path = ...
|
|
reloaded_keras_layer = hub.KerasLayer(gs_model_path)
|
|
inputs = tf.keras.layers.Input(
|
|
name="tweet__core__tweet__text", shape=(1,), dtype=tf.string
|
|
)
|
|
output = reloaded_keras_layer(inputs)
|
|
v7_model = tf.keras.models.Model(inputs=inputs, outputs=output)
|
|
pr_auc = tf.keras.metrics.AUC(curve="PR", name="pr_auc")
|
|
roc_auc = tf.keras.metrics.AUC(curve="ROC", name="roc_auc")
|
|
v7_model.compile(metrics=[pr_auc, roc_auc])
|
|
|
|
model.load_weights("...")
|
|
candidate_model = model
|
|
|
|
with mirrored_strategy.scope():
|
|
candidate_eval = candidate_model.evaluate(val_ds)
|
|
|
|
test_query = f"""
|
|
SELECT
|
|
{",".join(ptos_prototype.feature_names())},
|
|
has_media,
|
|
precision_nsfw,
|
|
{",".join(labels)},
|
|
...
|
|
"""
|
|
|
|
test = cq_loader.load_features(ptos_prototype, "", "", custom_query=test_query)
|
|
test = test.to_tf_dataset().map(parse_labeled_data)
|
|
|
|
print(test)
|
|
|
|
test_only_media = test.filter(lambda x, y: tf.equal(x["has_media"], True))
|
|
test_only_nsfw = test.filter(lambda x, y: tf.greater_equal(x["precision_nsfw"], 0.95))
|
|
test_no_media = test.filter(lambda x, y: tf.equal(x["has_media"], False))
|
|
test_media_not_nsfw = test.filter(
|
|
lambda x, y: tf.logical_and(
|
|
tf.equal(x["has_media"], True), tf.less(x["precision_nsfw"], 0.95)
|
|
)
|
|
)
|
|
for d in [test, test_only_media, test_only_nsfw, test_no_media, test_media_not_nsfw]:
|
|
print(d.reduce(0, lambda x, _: x + 1).numpy())
|
|
|
|
from dataclasses import asdict
|
|
|
|
from notebook_eval_utils import EvalConfig, SparseMultilabelEvaluator
|
|
|
|
|
|
def display_metrics(probs, targets, labels=labels):
|
|
eval_config = EvalConfig(prediction_threshold=0.5, precision_k=0.9)
|
|
for eval_mode, y_mask in [("implicit", np.ones(targets.shape))]:
|
|
print("Evaluation mode", eval_mode)
|
|
metrics = SparseMultilabelEvaluator.evaluate(
|
|
targets, np.array(probs), y_mask, classes=labels, eval_config=eval_config
|
|
)
|
|
metrics_df = pd.DataFrame.from_dict(
|
|
asdict(metrics)["per_topic_metrics"]
|
|
).transpose()
|
|
metrics_df["pos_to_neg"] = metrics_df["num_pos_samples"] / (
|
|
metrics_df["num_neg_samples"] + 1
|
|
)
|
|
display(metrics_df.median())
|
|
display(metrics_df)
|
|
return metrics_df
|
|
|
|
|
|
def eval_model(model, df):
|
|
with mirrored_strategy.scope():
|
|
targets = np.stack(list(df.map(lambda x, y: y).as_numpy_iterator()), axis=0)
|
|
df = df.padded_batch(BATCH_SIZE)
|
|
preds = model.predict(df)
|
|
return display_metrics(preds, targets)
|
|
|
|
|
|
subsets = {
|
|
"test": test,
|
|
"test_only_media": test_only_media,
|
|
"test_only_nsfw": test_only_nsfw,
|
|
"test_no_media": test_no_media,
|
|
"test_media_not_nsfw": test_media_not_nsfw,
|
|
}
|
|
|
|
metrics = {}
|
|
for name, df in subsets.items():
|
|
metrics[name] = eval_model(candidate_model, df)
|
|
[(name, m.pr_auc) for name, m in metrics.items()]
|
|
for name, x in [
|
|
(name, m.pr_auc.to_string(index=False).strip().split("\n"))
|
|
for name, m in metrics.items()
|
|
]:
|
|
print(name)
|
|
for y in x:
|
|
print(y.strip(), end="\t")
|
|
print(".")
|
|
for d in [test, test_only_media, test_only_nsfw, test_no_media, test_media_not_nsfw]:
|
|
print(d.reduce(0, lambda x, _: x + 1).numpy())
|