From 5a56b8d710d015a0c69c3f9e15c74c107366b977 Mon Sep 17 00:00:00 2001 From: DhruvaBansal00 Date: Wed, 25 Sep 2024 14:51:58 -0700 Subject: [PATCH] skip adding start/end tuple if key not found --- src/autolabel/confidence.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/autolabel/confidence.py b/src/autolabel/confidence.py index 44715e1c..d0f2fefb 100644 --- a/src/autolabel/confidence.py +++ b/src/autolabel/confidence.py @@ -192,9 +192,10 @@ def logprob_average_per_key( # We did not find the key in the logprobs so we set its confidence as 0 # This should not be possible if the LLM followed its guidelines. logprob_per_key[key] = 0 - start_token = mapping[loc] - end_token = mapping[loc + len(key_to_find) - 1] - locations.append((start_token, end_token, key)) + else: + start_token = mapping[loc] + end_token = mapping[loc + len(key_to_find) - 1] + locations.append((start_token, end_token, key)) locations.sort() # Here, the locations consist of the start and end *token* indices for each key # i.e for the keys A and B, we find the start and end tokens where they are found in the logprobs list