from vllm import LLM, SamplingParams
import json 

model_tokenizer_path = "Qwen2.5-Coder-14B-Instruct"
llm = LLM(
    model=model_tokenizer_path,
    tokenizer=model_tokenizer_path,
    # dtype=args.dtype,
    enforce_eager=True,
    max_model_len=30000,
    disable_custom_all_reduce=True,
    enable_prefix_caching=False,
    trust_remote_code=True,
    gpu_memory_utilization=0.98
)
sampling_params = SamplingParams(
    n=10,
    max_tokens=4000,
    temperature=0.2,
    top_p=0.95,
    frequency_penalty=0,
    presence_penalty=0,
    stop=["###"],
)


dataset = []
with open("/ainative/codefuse/user/448207/code_bench/livecodebench/code_generation_lite/test4.jsonl", "r") as f:
    for line in f:
        dataset.append(json.loads(line))

new_dataset = []
for item in dataset:
    question_content = item["question_content"]
    starter_code = item["starter_code"]
    new_dataset.append({'question_content': question_content, 'starter_code':starter_code})

# print(new_dataset[0])

def format_prompt(question):
    FORMATTING_MESSAGE_WITH_STARTER_CODE = "You will use the following starter code to write the solution to the problem and enclose your code within delimiters."
    FORMATTING_WITHOUT_STARTER_CODE = "You will read the inputs from stdin solve the problem and write the answer to stdout (do not directly test on the sample inputs). Enclose your code within delimiters."
    if question['starter_code']:
        prompt = f"### Instruction:\n{question['question_content']}\n{FORMATTING_MESSAGE_WITH_STARTER_CODE}\n```python\n{question['starter_code']}\n```\n\n### Response:\n" 
    else:
        prompt = f"### Instruction:\n{question['question_content']}\n{FORMATTING_WITHOUT_STARTER_CODE}\n\n### Response:\n"
        # ```python\n# YOUR CODE HERE\n```\n
    return prompt


for item in new_dataset:
    prompt = format_prompt(item)
    # prompt = item['question_content']
    print(f"prompt: {prompt}")

    vllm_outputs = llm.generate([prompt], sampling_params)
    output = [o.text for o in vllm_outputs[0].outputs]
    print(f"output: {output}")
    assert False