Print the number of tokens generated

helunwencser · helunwencser · commit 82f87137e8cb · 2024-11-11T14:38:49.000-08:00
This is useful for verifying the correctness of AttentionSink. Differential Revision: [D65784095](https://our.internmc.facebook.com/intern/diff/D65784095/) [ghstack-poisoned]
diff --git a/examples/models/llama/runner/eager.py b/examples/models/llama/runner/eager.py
@@ -91,10 +91,11 @@ def main() -> None:
         else runner.text_completion(
             prompt=args.prompt,
             temperature=args.temperature,
+            echo=True,
         )
     )
     if args.show_tokens:
-        print(f"Tokens: {generated_tokens}")
+        print(f"Generated {len(generated_tokens)} tokens: {generated_tokens}")
 
 
 if __name__ == "__main__":

Original file line number	Diff line number	Diff line change
`@@ -91,10 +91,11 @@ def main() -> None:`
`91`	`91`	`else runner.text_completion(`
`92`	`92`	`prompt=args.prompt,`
`93`	`93`	`temperature=args.temperature,`
	`94`	`+ echo=True,`
`94`	`95`	`)`
`95`	`96`	`)`
`96`	`97`	`if args.show_tokens:`
`97`		`- print(f"Tokens: {generated_tokens}")`
	`98`	`+ print(f"Generated {len(generated_tokens)} tokens: {generated_tokens}")`
`98`	`99`
`99`	`100`
`100`	`101`	`if __name__ == "__main__":`