Skip to content

Commit

Permalink
Add generation params
Browse files Browse the repository at this point in the history
Signed-off-by: Martín Santillán Cooper <msantillancooper@ibm.com>
  • Loading branch information
martinscooper committed Feb 5, 2025
1 parent 587bfc0 commit 84e15e9
Showing 1 changed file with 11 additions and 1 deletion.
12 changes: 11 additions & 1 deletion src/unitxt/metrics.py
Original file line number Diff line number Diff line change
Expand Up @@ -6016,7 +6016,17 @@ def compute(self, references: List[Any], prediction: Any, task_data: Dict) -> di

result = self.model.generate(
prompt=[processed_input],
params=self.generation_params,
params={
"decoding_method": "greedy",
"max_new_tokens": 20,
"temperature": 0,
"return_options": {
"token_logprobs": True,
"generated_tokens": True,
"input_text": True,
"top_n_tokens": 5,
},
},
)
generated_tokens_list = result[0]["results"][0]["generated_tokens"]
label, prob_of_risk = self.parse_output(generated_tokens_list)
Expand Down

0 comments on commit 84e15e9

Please sign in to comment.