Spaces:
Runtime error
Runtime error
# Generates positive movie reviews by tuning a pretrained model on IMDB dataset | |
# with a sentiment reward function | |
import json | |
import os | |
import sys | |
from typing import List | |
import torch | |
from datasets import load_dataset | |
from transformers import pipeline | |
import trlx | |
from trlx.data.default_configs import TRLConfig, default_ppo_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 get_negative_score(scores): | |
return dict(map(lambda x: tuple(x.values()), scores))["NEGATIVE"] | |
def main(hparams={}): | |
# Merge sweep config with default config if given | |
config = TRLConfig.update(default_ppo_config().to_dict(), hparams) | |
if torch.cuda.is_available(): | |
device = int(os.environ.get("LOCAL_RANK", 0)) | |
else: | |
device = -1 | |
sentiment_fn = pipeline( | |
"sentiment-analysis", | |
"lvwerra/distilbert-imdb", | |
top_k=2, | |
truncation=True, | |
batch_size=256, | |
device=device, | |
) | |
def dense_reward_fn(samples: List[str], prompts: List[str], outputs: List[str], tokenizer, **kwargs) -> List[float]: | |
# Reward positively for initially negative then positive review | |
# Reward functions should never receive padded text except for a single EOS at the end | |
# Reward function should return token rewards for just the response | |
first_halves = [".".join(sample.split(".")[: len(sample.split(".")) // 2]) for sample in samples] | |
negative_first_halves = list(map(get_negative_score, sentiment_fn(first_halves))) | |
second_halves = [".".join(sample.split(".")[len(sample.split(".")) // 2 :]) for sample in samples] | |
positive_second_halves = list(map(get_positive_score, sentiment_fn(second_halves))) | |
text_scores = [[f, s] for f, s in zip(negative_first_halves, positive_second_halves)] | |
tok_scores = [] | |
for sample, prompt, response, text_score in zip(samples, prompts, outputs, text_scores): | |
toks = tokenizer(response).input_ids | |
tok_score = [0] * len(toks) | |
tok_score[len(tok_score) // 2] = text_score[0] | |
tok_score[-1] = text_score[1] | |
tok_scores.append(tok_score) | |
return tok_scores | |
# Take few words off of movies reviews as prompts | |
imdb = load_dataset("imdb", split="train+test") | |
prompts = [" ".join(review.split()[:4]) for review in imdb["text"]] | |
trlx.train( | |
reward_fn=dense_reward_fn, | |
prompts=prompts, | |
eval_prompts=["I don't know much about Hungarian underground"] * 256, | |
config=config, | |
) | |
if __name__ == "__main__": | |
hparams = {} if len(sys.argv) == 1 else json.loads(sys.argv[1]) | |
main(hparams) | |