commit-0
diff --git a/‎examples/star/inference.py‎
Lines changed: 48 additions & 0 deletions b/‎examples/star/inference.py‎
Lines changed: 48 additions & 0 deletions
diff --git a/‎examples/star/star.py‎
Lines changed: 19 additions & 0 deletions b/‎examples/star/star.py‎
Lines changed: 19 additions & 0 deletions
@@ -0,0 +1,48 @@
+from typing import List
+from datasets import Dataset
+from vllm import LLM, SamplingParams
+
+def generate_predictions(
+    model_name: str,
+    dataset: Dataset,
+    temperature: float = 1.0,
+    n: int = 1
+) -> List[List[str]]:
+    """
+    Generate predictions for a given dataset using a specified language model and
+    sampling parameters. The function loads the dataset, constructs prompts from
+    each example, and obtains generated predictions. The resulting predictions are
+    then added as a new column to the dataset.
+
+    Args:
+        model_name (str): Name of the model to use for generation.
+        dataset (Dataset): The Dataset object.
+        temperature (float, optional): Temperature setting for the model's
+            sampling strategy. Default is 1.0.
+        n (int, optional): Number of sampling runs per prompt. Default is 1.
+
+    Returns:
+        predictions (List[List[str]]): Predictions on the dataset.
+    """
+    sampling_params = SamplingParams(n=n, temperature=temperature, max_tokens=512)
+    llm = LLM(model=model_name)
+
+    prompts: List[str] = []
+    for example in dataset:
+        prompt = (
+            f"{example['text']} Your code should satisfy these tests:\n\n"
+            f"{'\n'.join(example['test_list'])}"
+        )
+        prompts.append(prompt)
+
+    outputs = llm.generate(prompts, sampling_params)
+
+    results: List[List[str]] = []
+    for output in outputs:
+        generated_texts = [one.text for one in output.outputs]
+        results.append(generated_texts)
+    return results
+    #out_name = dataset_name.split("/")[-1]
+    #out_name = f"wentingzhao/{out_name}_predictions_{n}"
+    #ds.push_to_hub(out_name)
+
@@ -0,0 +1,19 @@
+"""Main STaR Loop"""
+import argparse
+from datasets import Dataset, load_dataset
+from inference import generate_predictions
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name", type=str, required=True, help="model to use")
+    parser.add_argument("--dataset_name", type=str, required=True, help="dataset to use")
+    parser.add_argument("--temperature", type=float, default=1)
+    parser.add_argument("-n", type=int, default=1)
+    args = parser.parse_args()
+
+    ds = load_dataset(args.dataset_name)
+    assert "train" in ds
+    samples = generate_predictions(args.model_name, ds["train"], args.temperature, args.n)
+
+if __name__ == '__main__':
+    main()