added pipeline

rewardModel.py

@@ -0,0 +1,113 @@
+import torch
+
+from datasets import load_dataset
+from transformers import pipeline, AutoTokenizer
+
+from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead
+from trl.core import LengthSampler
+
+
+dataset = load_dataset("imdb")
+# print(dataset['train']['text'][:5])
+# print(dataset['train']['label'][:5])x
+
+def build_dataset(config, dataset_name='imdb', input_min = 2, input_max=8):
+    
+
+    tokenizer = AutoTokenizer.from_pretrained(config.model_name)
+    tokenizer.pad_token = tokenizer.eos_token
+
+
+    ds = load_dataset(dataset_name, split="train")
+    ds = ds.rename_columns({"text": "review"})
+    ds = ds.filter(lambda x: len(x['review']) > 200, batched=False)
+
+    input_size = LengthSampler(input_min, input_max)
+
+    def tokenize(sample):
+        sample['input_ids'] = tokenizer.encode(sample['review'])[: input_size()]
+        sample['query'] = tokenizer.decode(sample["input_ids"])
+        return sample
+    
+    ds = ds.map(tokenize, batched=False)
+    ds.set_format(type="torch")
+    return ds
+
+
+def collator(data):
+    return dict((key, [d[key] for d in data]) for key in data[0])
+    
+
+
+if __name__ == '__main__':   #START
+    print("running")
+    config = PPOConfig(
+        model_name="lvwerra/gpt2-imdb",
+        learning_rate=1.41e-5,
+        log_with="wandb"
+    )
+
+    import wandb
+    wandb.init()
+
+    dataset = build_dataset(config)
+    model = AutoModelForCausalLMWithValueHead.from_pretrained(config.model_name)
+
+    ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(config.model_name)  # this is the frozen model which we'll compare to
+
+    tokenizer = AutoTokenizer.from_pretrained(config.model_name)
+    tokenizer.pad_token = tokenizer.eos_token
+
+    ppo_trainer = PPOTrainer(config,  model, ref_model, tokenizer, dataset, data_collator=collator)    
+    device = ppo_trainer.accelerator.device    
+    if ppo_trainer.accelerator.num_processes == 1:
+        device = 0 if torch.cuda.is_available() else "cpu"
+
+    sentiment_pipe = pipeline("sentiment-analysis", model='lvwerra/distilbert-imdb', device=device)
+
+    sent_kwargs = {'return_all_scores': True, "function_to_apply": "none", "batch_size": 16}
+    text = "this was really bad!!"
+    print(sentiment_pipe(text, **sent_kwargs))
+
+    text = "this was really good!!"
+    print(sentiment_pipe(text, **sent_kwargs))
+
+    output_min_length = 4
+    output_max_length = 16
+    output_length_sampler = LengthSampler(output_min_length, output_max_length)
+
+    response_generation_kwargs = {
+        "min_length": -1,
+        "top_k": 0.0,
+        "top_p": 1.0,
+        "do_sample": True,
+        "pad_token_id": tokenizer.eos_token_id,
+
+    }
+
+    for epoch, batch in enumerate(ppo_trainer.dataloader):
+        
+        query_tensors = batch["input_ids"]
+
+        response_tensors = []
+
+        for query in query_tensors:
+            gen_len = output_length_sampler()
+            response_generation_kwargs["max_new_tokens"] = gen_len
+            response = ppo_trainer.generate(query, **response_generation_kwargs)
+            response_tensors.append(response.squeeze()[-gen_len:]) # only take the piece generated, remove acutal prompt
+        batch['response'] = [tokenizer.decode(r.squeeze()) for r in response_tensors]
+
+        texts = [q + r for q, r in zip(batch['query'], batch['response'])]
+
+        pipe_outputs = sentiment_pipe(texts, **sent_kwargs)  # this gives a dictionary, we want the positive, which is our rewards 
+
+        rewards = [torch.tensor(output[1]["score"]) for output in pipe_outputs]
+
+        stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
+
+        ppo_trainer.log_stats(stats, batch, rewards)
+
+    # model.save_pretrained('gpt-imbd-pos-v2', push_to_hub=False)
+
+    # tokenizer.save_pretrained('gpt-imbd-pos-v2', push_to_hub=False)