Merge pull request #23 from aalok-sathe/bugfix/cli-propogate-bos-toke…

…n-arg CLI wasn't actually using the BOS token arg; make sure to pass it in case of gpt-like models
aalok-sathe · Dec 3, 2024 · 2c6f250 · 2c6f250
2 parents bd49e84 + f3abb81
commit 2c6f250
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/surprisal/__main__.py b/surprisal/__main__.py
@@ -57,8 +57,11 @@
         args.model_name_or_path, model_class=args.model_class
     )
 
-    [surp] = m.surprise(args.text)
-
+    if args.model_class in ('gpt', 'causal'): 
+        [surp] = m.surprise(args.text, use_bos_token=not args.no_bos_token)
+    else:
+        [surp] = m.surprise(args.text)
+
     if args.plot:
         from matplotlib import pyplot as plt