@@ -6,7 +6,7 @@
 from trlx.data.configs import TRLConfig
 
 
-def reward_fn(samples):
+def reward_fn(samples, **kwargs):
     "Gives a negative count of rooms for each sample"
     return [-sample.count(":") for sample in samples]
 
 
@@ -29,7 +29,7 @@ def main(hparams={}):
         device=0 if int(os.environ.get("LOCAL_RANK", 0)) == 0 else -1,
     )
 
-    def metric_fn(samples: List[str]) -> Dict[str, List[float]]:
+    def metric_fn(samples: List[str], **kwargs) -> Dict[str, List[float]]:
         sentiments = list(map(get_positive_score, sentiment_fn(samples)))
         return {"sentiments": sentiments}
 
 
@@ -38,7 +38,7 @@ def main(hparams={}):
         device=device,
     )
 
-    def reward_fn(samples: List[str]) -> List[float]:
+    def reward_fn(samples: List[str], **kwargs) -> List[float]:
         sentiments = list(map(get_positive_score, sentiment_fn(samples)))
         return sentiments
 
 
@@ -44,6 +44,6 @@ method:
   two_qs: true
   gen_kwargs:
     max_new_tokens: 9
-    top_k: 1
-    beta: 100
+    top_k: 10
+    beta: [0, 1, 100]
     temperature: 1.0
@@ -23,8 +23,9 @@ def main(hparams={}):
         GPT2Config(n_layer=6, n_embd=144, vocab_size=23),
         dataset=(walks, rewards),
         eval_prompts=eval_prompts,
-        metric_fn=metric_fn,
+        metric_fn=lambda samples, **kwargs: metric_fn(samples),
         config=config,
+        stop_sequences=["|"],
     )
 
 
 
@@ -17,10 +17,10 @@ def main(hparams={}):
 
     trlx.train(
         "CarperAI/randomwalks",
-        reward_fn=lambda walks: metric_fn(walks)["optimality"],
+        reward_fn=lambda samples, **kwargs: metric_fn(samples)["optimality"],
         prompts=prompts,
         eval_prompts=prompts,
-        metric_fn=metric_fn,
+        metric_fn=lambda samples, **kwargs: metric_fn(samples),
         config=config,
     )
 
 
@@ -25,16 +25,14 @@
 
 if __name__ == "__main__":
 
-    def reward_fn(samples: List[str]):
-        sep_token = tokenizer.sep_token
-        articles = [sample.split(sep_token)[0].strip() for sample in samples]
-        predicted_summaries = [sample.split(sep_token)[1].strip() for sample in samples]
-        labels = [prompt_label[sample] for sample in articles]
+    def reward_fn(samples: List[str], prompts: List[str], outputs: List[str]):
+        original_summaries = [prompt_label[prompt.strip()] for prompt in prompts]
         scores = [
-            meteor.compute(predictions=[summary], references=[label])
-            for (summary, label) in zip(predicted_summaries, labels)
+            meteor.compute(predictions=[output.strip()], references=[original])[
+                "meteor"
+            ]
+            for (original, output) in zip(original_summaries, outputs)
         ]
-        scores = [score["meteor"] for score in scores]
         return scores
 
     dataset = load_dataset("cnn_dailymail", "3.0.0", cache_dir="data")
 
@@ -19,6 +19,7 @@ install_requires =
     torchtyping
     transformers>=4.21.2
     tqdm
+    rich
     wandb>=0.13.5
     ray>=2.0.1
     tabulate>=0.9.0
 
@@ -1,8 +1,8 @@
 from time import time
-from typing import Callable, Optional
 
 import ray
 import torch
+import torch.nn.functional as F
 
 from trlx.data.accelerate_base_datatypes import PromptBatch
 from trlx.data.ppo_types import PPORLElement
@@ -24,8 +24,6 @@ def __init__(
         self,
         trainer: BaseRLTrainer,
         pipeline: BasePipeline,
-        reward_fn: Callable,
-        metric_fn: Optional[Callable] = None,
         chunk_size: int = 512,
     ):
         self.pipeline = pipeline
@@ -43,8 +41,6 @@ def __init__(
             self.ref_model.to(self.trainer.accelerator.device)
 
         self.trainer.orch = self
-        self.trainer.reward_fn = reward_fn
-        self.trainer.metric_fn = metric_fn
 
         self.running = RunningMoments()
         self.ref_mean = self.trainer.config.method.ref_mean
@@ -65,9 +61,6 @@ def make_experience(self, num_rollouts: int = 1024, iter_count: int = 0):  # noq
         stats = {}
         clock = Clock()
         while len(ppo_rl_elements) < num_rollouts:
-            if self.trainer.accelerator.is_main_process:
-                print(f"Making experience {len(ppo_rl_elements)} / {num_rollouts}")
-
             # Get next batch in prompt dataset and refresh if exhausted
             try:
                 batch: PromptBatch = next(self.pipeline_iterator)
@@ -79,30 +72,38 @@ def make_experience(self, num_rollouts: int = 1024, iter_count: int = 0):  # noq
             samples = self.trainer.generate(**batch)
             stats["time/exp_generate"] = time() - exp_generate_time
 
-            if self.trainer.config.model.model_arch_type == "seq2seq":
-                response_tensors = samples
-            else:
-                query_tensors = batch.input_ids
-                response_tensors = samples[:, query_tensors.shape[1] :]
-
-            texts = self.trainer.tokenizer.batch_decode(
-                samples, skip_special_tokens=True
+            query_tensors = batch.input_ids
+            device = samples.device
+            str_samples, str_prompts, str_outputs = self.trainer.decode(
+                query_tensors, samples
             )
 
-            if self.trainer.config.model.model_arch_type == "seq2seq":
-                articles = self.trainer.tokenizer.batch_decode(
-                    batch.input_ids, skip_special_tokens=True
+            # Convert trimmed samples back into tensors for another head pass
+            # This can be defered, instead letting the pass to made over the original samples
+            # after unbinding and truncating operations lower are fixed
+            outputs = self.trainer.tokenizer(str_outputs).input_ids
+            outputs = list(map(torch.LongTensor, outputs))
+            maxsize = max(map(len, outputs))
+            outputs = [
+                F.pad(
+                    output,
+                    (0, maxsize - len(output)),
+                    value=self.trainer.tokenizer.pad_token_id,
                 )
-                sep_token = self.trainer.tokenizer.sep_token
-                texts = [
-                    f"{article}{sep_token}{response}"
-                    for article, response in zip(articles, texts)
-                ]
+                for output in outputs
+            ]
+            response_tensors = torch.vstack(outputs).to(device)
 
             exp_score_time = time()
+
             scores = torch.tensor(
-                self.score(texts), device=samples.device, dtype=torch.float
-            )
+                self.trainer.reward_fn(
+                    samples=str_samples,
+                    prompts=str_prompts,
+                    outputs=str_outputs,
+                ),
+                dtype=float,
+            ).to(device)
             stats["time/exp_score"] = time() - exp_score_time
 
             # store statistics of the initial rollout as reference
@@ -125,9 +126,8 @@ def make_experience(self, num_rollouts: int = 1024, iter_count: int = 0):  # noq
 
             # Precompute logprobs, values
             if self.trainer.config.model.model_arch_type == "seq2seq":
-                response_tensors = response_tensors
-                attention_mask = batch.attention_mask.to(response_tensors.device)
-                query_tensors = batch.input_ids.to(response_tensors.device)
+                attention_mask = batch.attention_mask.to(device)
+                query_tensors = batch.input_ids.to(device)
                 with torch.no_grad():
                     outputs = self.trainer.model(
                         input_ids=query_tensors,
@@ -150,12 +150,12 @@ def make_experience(self, num_rollouts: int = 1024, iter_count: int = 0):  # noq
                         ).logits
             else:
                 all_tokens = torch.cat(
-                    (query_tensors.to(response_tensors.device), response_tensors), dim=1
+                    (query_tensors.to(device), response_tensors), dim=1
                 )
                 attention_mask = (
                     all_tokens.not_equal(self.trainer.tokenizer.pad_token_id)
                     .long()
-                    .to(all_tokens.device)
+                    .to(device)
                 )
                 with torch.no_grad():
                     logits, *_, values = self.trainer.model(
@@ -175,7 +175,7 @@ def make_experience(self, num_rollouts: int = 1024, iter_count: int = 0):  # noq
                             attention_mask=attention_mask,
                             return_dict=False,
                         )
-                        ref_logits = ref_logits.to(self.trainer.accelerator.device)
+                        ref_logits = ref_logits.to(device)
 
             if self.trainer.config.model.model_arch_type == "seq2seq":
                 logprobs = logprobs_from_logits(
 
@@ -7,6 +7,8 @@
 
 import wandb
 
+from trlx.utils import significant
+
 import wandb.apis.reports as wb  # isort: skip
 
 
@@ -39,10 +41,6 @@ def parse_result(result):
     return out
 
 
-def significant(x):
-    return round(x, 1 - int(math.floor(math.log10(x))))
-
-
 def log_trials(trial_path: str, project_name: str):
     trial_path = Path(trial_path)
     files = os.listdir(trial_path)
Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@ def main(hparams={}):`
`29`	`29`	`device=0 if int(os.environ.get("LOCAL_RANK", 0)) == 0 else -1,`
`30`	`30`	`)`
`31`	`31`
`32`		`- def metric_fn(samples: List[str]) -> Dict[str, List[float]]:`
	`32`	`+ def metric_fn(samples: List[str], **kwargs) -> Dict[str, List[float]]:`
`33`	`33`	`sentiments = list(map(get_positive_score, sentiment_fn(samples)))`
`34`	`34`	`return {"sentiments": sentiments}`
`35`	`35`
Original file line number	Diff line number	Diff line change
`@@ -38,7 +38,7 @@ def main(hparams={}):`
`38`	`38`	`device=device,`
`39`	`39`	`)`
`40`	`40`
`41`		`- def reward_fn(samples: List[str]) -> List[float]:`
	`41`	`+ def reward_fn(samples: List[str], **kwargs) -> List[float]:`
`42`	`42`	`sentiments = list(map(get_positive_score, sentiment_fn(samples)))`
`43`	`43`	`return sentiments`
`44`	`44`
Original file line number	Diff line number	Diff line change
`@@ -23,8 +23,9 @@ def main(hparams={}):`
`23`	`23`	`GPT2Config(n_layer=6, n_embd=144, vocab_size=23),`
`24`	`24`	`dataset=(walks, rewards),`
`25`	`25`	`eval_prompts=eval_prompts,`
`26`		`- metric_fn=metric_fn,`
	`26`	`+ metric_fn=lambda samples, **kwargs: metric_fn(samples),`
`27`	`27`	`config=config,`
	`28`	`+ stop_sequences=["\|"],`
`28`	`29`	`)`
`29`	`30`
`30`	`31`