add top_p

Signed-off-by: SumanthRH <[email protected]>

skythought/skythought_evals/inference_and_check.py

@@ -92,7 +92,12 @@ def fetch_responses_ray(conversations, max_tokens, temp, args):
         ds = ds.repartition(num_partitions=num_replicas)
     workload = EvalWorkload(
         dataset=ds,
-        sampling_params={"n": args.n, "max_tokens": max_tokens, "temperature": temp},
+        sampling_params={
+            "n": args.n,
+            "max_tokens": max_tokens,
+            "temperature": temp,
+            "top_p": args.top_p,
+        },
     )
     pipeline = Pipeline(
         engine_cfg,
@@ -138,7 +143,7 @@ def inference(llm, conversations, max_tokens, temp, args):
         responses = [Response.from_openai_response(response) for response in responses]
     else:
         sampling_params = SamplingParams(
-            max_tokens=max_tokens, temperature=temp, n=args.n
+            max_tokens=max_tokens, temperature=temp, n=args.n, top_p=args.top_p
         )
         responses = llm.chat(
             messages=conversations, sampling_params=sampling_params, use_tqdm=True
@@ -629,6 +634,12 @@ def main():
         "'auto' refers to automatically inferring dtype for the model",
         default="float32",
     )
+    parser.add_argument(
+        "--top_p",
+        type=float,
+        default=1,
+        help="Sampling parameter `top_p`",
+    )
     args = parser.parse_args()
     # load ray config
     if args.use_ray:
@@ -660,6 +671,12 @@ def main():
 
     temperatures = [1] if args.model.startswith("openai/o1") else args.temperatures
 
+    if args.top_p < 1 and args.model.startswith("openai/o1"):
+        print(
+            "OpenAI o1 models do not support `top_p` sampling. Resetting `top_p` to 1"
+        )
+        args.top_p = 1
+
     print(f"Temperature: {temperatures}")
     max_tokens = args.max_tokens
     if temperatures == [0] and args.n > 1:

skythought/skythought_evals/util/metrics.py

@@ -13,7 +13,7 @@ def _pass_at_k(n, c, k):
     """
     if n - c < k:
         return 1.0
-    return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1))
+    return float(1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1)))
 
 
 def pass_at_k(N: int, temp_to_scores: Dict[str, Dict[str, Any]]):