Spaces:
Runtime error
Runtime error
File size: 1,666 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 |
from typing import Dict, List
from datasets import load_dataset
from transformers import pipeline
import trlx
from trlx.data.default_configs import default_ilql_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"]
default_config = default_ilql_config()
def main(hparams={}):
# Merge sweep config with default config if given
config = default_config.evolve(
train=dict(
seq_length=1024,
batch_size=512,
total_steps=200,
trainer="NeMoILQLTrainer",
trainer_kwargs=dict(
pretrained_model=None,
megatron_cfg="megatron_20b.yaml",
),
),
method=dict(
gen_kwargs=dict(
beta=2.0,
temperature=0.9,
)
),
)
config = config.evolve(**hparams)
sentiment_fn = pipeline(
"sentiment-analysis",
"lvwerra/distilbert-imdb",
top_k=2,
truncation=True,
batch_size=256,
device=-1,
)
def metric_fn(samples: List[str], **kwargs) -> Dict[str, List[float]]:
sentiments = list(map(get_positive_score, sentiment_fn(samples)))
return {"sentiments": sentiments}
imdb = load_dataset("imdb", split="train+test")
trlx.train(
samples=imdb["text"],
rewards=imdb["label"],
eval_prompts=["I don't know much about Hungarian underground"] * 128,
metric_fn=metric_fn,
config=config,
)
if __name__ == "__main__":
main()
|