Spaces:
Runtime error
Runtime error
File size: 3,529 Bytes
fa6856c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import json
import os
from argparse import ArgumentParser
from typing import Dict, List
from datasets import load_dataset
from transformers import pipeline
import trlx
from trlx.data.default_configs import TRLConfig, default_sft_config
def get_positive_score(scores):
"Extract value associated with a positive sentiment from pipeline's output"
return dict(map(lambda x: tuple(x.values()), scores))["POSITIVE"]
def preprocess(instruction: str, input: str, output: str):
"""Build Alpaca prompt and output from instruction and input/output examples"""
if input:
prefix = (
"Below is an instruction that describes a task, paired with an input that provides further context. "
"Write a response that appropriately completes the request."
)
prompt = f"{prefix}\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n"
return [prompt, output]
else:
prefix = (
"Below is an instruction that describes a task. Write a response that appropriately completes the request."
)
prompt = f"{prefix}\n\n### Instruction:\n{instruction}\n\n### Response:\n"
return [prompt, output]
def main(hparams={}, model_name="EleutherAI/gpt-j-6B", dataset="tatsu-lab/alpaca"):
config = default_sft_config()
config = config.evolve(
train=dict(
total_steps=2400,
batch_size=4,
seq_length=1024,
),
model=dict(
model_path=model_name,
),
tokenizer=dict(
tokenizer_path=model_name,
),
optimizer=dict(kwargs=dict(lr=2e-5)),
scheduler=dict(kwargs=dict(eta_min=2e-5)),
method=dict(
gen_kwargs=dict(
max_new_tokens=256,
)
),
)
# Merge sweep config with default config if given
config = TRLConfig.update(config.to_dict(), hparams)
# alpaca = load_dataset("tatsu-lab/alpaca", split="train")
alpaca = load_dataset(dataset, split="train")
alpaca = [preprocess(x["instruction"], x["input"], x["output"]) for x in alpaca]
sentiment_fn = pipeline(
"sentiment-analysis",
"lvwerra/distilbert-imdb",
top_k=2,
truncation=True,
batch_size=256,
device=0 if int(os.environ.get("LOCAL_RANK", 0)) == 0 else -1,
)
def metric_fn(samples: List[str], prompts: List[str], outputs: List[str]) -> Dict[str, List[float]]:
sentiments = list(map(get_positive_score, sentiment_fn(outputs)))
return {"sentiments": sentiments}
imdb = load_dataset("imdb", split="test")
bad_reviews = imdb.filter(lambda sample: sample["label"] == 0).select(range(256))
zs_rewrite = [preprocess("Rewrite the input into a positive review.", x["text"][:1024], "")[0] for x in bad_reviews]
trainer = trlx.train(
samples=alpaca,
eval_prompts=zs_rewrite,
metric_fn=metric_fn,
config=config,
)
slug = f"{model_name.split('/')[-1]}-{dataset.split('/')[-1]}"
trainer.save_pretrained(f"{slug}-sft")
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("override_hparams", type=str, default="{}", nargs="?")
parser.add_argument("--model_name", type=str, default="EleutherAI/gpt-j-6B")
parser.add_argument("--dataset", type=str, default="tatsu-lab/alpaca")
args = parser.parse_args()
hparams = json.loads(args.override_hparams)
main(hparams, args.model_name, args.dataset)
|