Spaces:
Runtime error
Runtime error
import json | |
import os | |
import sys | |
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 main(hparams={}): | |
# Merge sweep config with default config if given | |
config = TRLConfig.update(default_sft_config().to_dict(), hparams) | |
imdb = load_dataset("imdb", split="train+test") | |
# Finetune on only positive reviews | |
imdb = imdb.filter(lambda sample: sample["label"] == 1) | |
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], **kwargs) -> Dict[str, List[float]]: | |
sentiments = list(map(get_positive_score, sentiment_fn(samples))) | |
return {"sentiments": sentiments} | |
trainer = trlx.train( | |
samples=imdb["text"], | |
eval_prompts=["I don't know much about Hungarian underground"] * 64, | |
metric_fn=metric_fn, | |
config=config, | |
) | |
trainer.save_pretrained("reviews-sft") | |
if __name__ == "__main__": | |
hparams = {} if len(sys.argv) == 1 else json.loads(sys.argv[1]) | |
main(hparams) | |