Refactor: Move return_hidden_states to the generate input (#3985)

Co-authored-by: Beichen-Ma <mabeichen12@gmail.com>
2025-03-01 20:51:29 -05:00
parent 18bb216c28
commit 40782f05d7
12 changed files with 54 additions and 44 deletions
--- a/examples/runtime/engine/hidden_states.py
+++ b/examples/runtime/engine/hidden_states.py
@@ -26,10 +26,11 @@ def main():
        "temperature": 0.8,
        "top_p": 0.95,
        "max_new_tokens": 10,
-        "return_hidden_states": True,
    }

-    outputs = llm.generate(prompts, sampling_params=sampling_params)
+    outputs = llm.generate(
+        prompts, sampling_params=sampling_params, return_hidden_states=True
+    )
    for prompt, output in zip(prompts, outputs):
        print("===============================")
        print(