fix: Handle EOT as well as __PAD__ segments for pooler output EOT calc

2023-11-18 23:04:06 -08:00
parent dba41a2c4d
commit 2fa1b95045
1 changed files with 1 additions and 1 deletions
@@ -219,7 +219,7 @@ class PrompLangCLIPTextTransformer(CLIPTextTransformer):
                if isinstance(seg_or_action, Action):
                    idx += seg_or_action.token_length()
                else:
-                    if seg_or_action.text == '__PAD__':
+                    if seg_or_action.text == "__PAD__" or seg_or_action.text == "[EOT]":
                        break

                    # Is a segment, and isn't the pad segment