wip

repepo/core/evaluate.py

@@ -255,7 +255,6 @@ def get_prediction(
     pipeline: Pipeline,
     example: Example,
     evaluators: Sequence[Evaluator] = [LogitDifferenceEvaluator()],
-    n_generation: int = 0,
     slim_results: bool = False,
 ) -> EvalPrediction:
     positive_probs = pipeline.calculate_output_logprobs(
@@ -272,10 +271,10 @@ def get_prediction(
     example_metrics = {}
     for evaluator in evaluators:
         example_metrics[evaluator.get_metric_name()] = evaluator.score_prediction(pred)
-    for _ in range(n_generation):
-        prompt = pipeline.build_generation_prompt(example.positive)
-        response = pipeline.generate(example.positive)
-        pred.generations.append(Completion(prompt=prompt, response=response))
+    # for _ in range(n_generation):
+    #     prompt = pipeline.build_generation_prompt(example.positive)
+    #     response = pipeline.generate(example.positive)
+    #     pred.generations.append(Completion(prompt=prompt, response=response))
     pred.metrics = example_metrics
     return pred
 

repepo/core/pipeline.py

@@ -59,7 +59,7 @@ def __repr__(self) -> str:
 
 @dataclass
 class PipelineContext:
-    method: Literal["generate", "logprobs"]
+    method: Literal["generate", "logprobs", "forward"]
     base_prompt: str
     full_prompt: str
     inputs: Any
@@ -118,6 +118,7 @@ def build_full_prompt(self, completion: Completion) -> str:
             self.formatter.format_conversation(completion, self.conversation_history)
         )
 
+    @torch.no_grad()
     def generate(
         self,
         completion: Completion,
@@ -147,12 +148,36 @@ def generate(
                 return outputs_str.replace(base_prompt, "")
             return outputs_str
         raise RuntimeError("Should never get here")
+
+    @torch.no_grad()
+    def __call__(self, text: str):
+        """A lightweight wrapper around model"""
+        # with ExitStack() as stack:
+        #     for hook in self.hooks:
+        #         stack.enter_context(hook(context))
+        # base_prompt = self.build_generation_prompt(completion)
+        # full_prompt = self.build_full_prompt(completion)
+        inputs: Any = self.tokenizer(text, return_tensors="pt")
+        inputs = inputs.to(self.model.device)
+        context = PipelineContext(
+            method="forward",
+            base_prompt="dummy",
+            full_prompt="dummy",
+            inputs=inputs,
+            pipeline=self,
+        )
+        with ExitStack() as stack:
+            for hook in self.hooks:
+                stack.enter_context(hook(context))
+            outputs = self.model(**inputs, output_hidden_states=False, return_dict=True)
+            return outputs
 
     @torch.no_grad()
     def calculate_output_logprobs(
         self, completion: Completion, slim_results: bool = False
     ) -> TextProbs:
         """Calculate the logprobs for each token in the prompt + output"""
+
         base_prompt = self.build_generation_prompt(completion)
         full_prompt = self.build_full_prompt(completion)
         inputs: Any = self.tokenizer(full_prompt, return_tensors="pt")