Naman-ntc · Apr 12, 2024 · Apr 12, 2024
diff --git a/lcb_runner/evaluation/compute_scores.py b/lcb_runner/evaluation/compute_scores.py
@@ -2,6 +2,7 @@
 import argparse
 from datetime import datetime
 
+from lcb_runner.lm_styles import LanguageModelStore
 from lcb_runner.utils.scenarios import Scenario
 from lcb_runner.utils.path_utils import get_eval_all_output_path
 
@@ -50,7 +51,8 @@ def get_parser():
     args = parser.parse_args()
 
     if args.eval_all_file is None:
-        args.eval_all_file = get_eval_all_output_path(args.model, args)
+        model = LanguageModelStore[args.model]
+        args.eval_all_file = get_eval_all_output_path(model, args)
 
     return args
 

diff --git a/lcb_runner/lm_styles.py b/lcb_runner/lm_styles.py
@@ -23,6 +23,9 @@ class LMStyle(Enum):
     MagiCoder = "MagiCoder"
     OC = "OC"
 
+    Qwen1point5 = "Qwen1point5"
+    Smaug2 = "Smaug2"
+
 
 @dataclass
 class LanguageModel:
@@ -282,6 +285,20 @@ def __hash__(self) -> int:
         datetime(2023, 9, 1),
         link="https://huggingface.co/stabilityai/stable-code-3b/",
     ),
+    LanguageModel(
+        "/abacus/models/Qwen1.5-72B-Chat/",
+        "Qwen-1.5-72B-Chat ",
+        LMStyle.Qwen1point5,
+        datetime(2024, 3, 31),
+        link="https://huggingface.co/qwen/Qwen1.5-72B-Chat/",
+    ),
+    LanguageModel(
+        "/abacus/models/Smaug-2-72B/",
+        "Smaug-2-72B ",
+        LMStyle.Smaug2,
+        datetime(2024, 3, 31),
+        link="https://huggingface.co/abacusai/Smaug-2-72B/",
+    ),
 ]
 
 LanguageModelStore: dict[str, LanguageModel] = {

diff --git a/lcb_runner/prompts/code_generation.py b/lcb_runner/prompts/code_generation.py
@@ -72,6 +72,30 @@ def get_deepseekcode_question_template_answer(question: CodeGenerationProblem):
     return prompt
 
 
+def get_qwen_question_template_answer(question: CodeGenerationProblem):
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained(
+        "/abacus/models/Qwen1.5-72B-Chat/",
+        padding_side="left",
+        use_fast=False
+    )
+    prompt = "You will be given a question (problem specification) and will generate a correct Python program that matches the specification and passes all tests. You will NOT return anything except for the program.\n\n"
+    prompt += f"Question:\n{question.question_content}\n\n"
+    if question.starter_code:
+        prompt += f"{PromptConstants.FORMATTING_MESSAGE_WITH_STARTER_CODE}\n"
+        prompt += f"```python\n{question.starter_code}\n```\n\n"
+    else:
+        prompt += f"{PromptConstants.FORMATTING_WITHOUT_STARTER_CODE}\n\n"
+        prompt += f"```python\n# YOUR CODE HERE\n```\n\n"
+
+    messages = [{'role': 'system', 'content': PromptConstants.SYSTEM_MESSAGE_GENERIC},
+                {'role': 'user', 'content': prompt}]
+
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True,
+                                           truncation=False, padding=False)
+    return prompt
+
+
 def get_magicoder_question_template_answer(question: CodeGenerationProblem):
     prompt = f"You will be given a question (problem specification) and will generate a correct Python program that matches the specification and passes all tests. You will NOT return anything except for the program.\n\n"
     prompt += f"Question:\n{question.question_content}\n\n"
@@ -242,6 +266,10 @@ def format_prompt_generation(
         prompt += f"{get_generic_question_template_answer(question)}"
         return prompt
 
+    if LanguageModelStyle == LMStyle.Smaug2 or LanguageModelStyle == LMStyle.Qwen1point5:
+        prompt = f"{get_qwen_question_template_answer(question)}"
+        return prompt
+
     if LanguageModelStyle in [
         LMStyle.DeepSeekBase,
         LMStyle.CodeLLaMaBase,

diff --git a/lcb_runner/runner/vllm_runner.py b/lcb_runner/runner/vllm_runner.py
@@ -27,6 +27,9 @@ def __init__(self, args, model):
             stop=self.args.stop,
         )
 
+    def _run_single(self, prompt: str) -> list[str]:
+        pass
+
     def run_batch(self, prompts: list[str]) -> list[list[str]]:
         outputs = [None for _ in prompts]
         remaining_prompts = []