I’m new to LLMs and trying to reproduce the experiments from "The Consensus Game: Language Model Generation via Equilibrium Search." However, I’m struggling to find the right prompt format to achieve the expected results from the paper. I’m attempting to get the probability of each option from the model’s response, but the performance is no better than random guessing (~19% for multiple choice with 4 options, ). I’m not sure what’s wrong with my code. Has anyone worked on this before and found a better approach? (I try to test it on MMLU dataset) Additionally, I noticed that Llama3-8B does not seem sensitive to "incorrect" and "correct" answers. Even when I explicitly ask it to provide an incorrect answer, the probabilities of the options remain nearly the same as when I ask for the correct answer. Here is my code def format_subject(subject): l = subject.split("_") s = "" for entry in l: s += " " + entry return s def build_generator_prompt( subject, target_question, target_choices, get_correct ): prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|> You are a helpful assistant. Answer the following multiple-choice question about {format_subject(subject)}. Provide only the letter corresponding to the answer.<|eot_id|> <|start_header_id|>user<|end_header_id|> Question: {target_question} Options: """ for i, c in enumerate(target_choices): prompt += "\n{}".format(c) if get_correct: prompt += "\nAnswer:<eot_id>\n<|start_header_id|>assistant<|end_header_id|>" else: prompt += "\nIncorrect Answer:<eot_id>\n<|start_header_id|>assistant<|end_header_id|>" print(prompt) return prompt def get_generator_answer_probs(model, tokenizer, prompt_text, choices_list): inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device) outputs = model(**inputs) logits = outputs.logits[0, -1] choice_token_ids = [] for choice in ["A", "B", "C", "D"]: token_id = tokenizer.encode(choice, add_special_tokens=False) if len(token_id) == 1: choice_token_ids.append(token_id[0]) else: raise ValueError(f"Choice '{choice}' is tokenized into multiple tokens: {token_id}") choice_logits = logits[torch.tensor(choice_token_ids).to(model.device)] probs = torch.softmax(choice_logits, dim=-1).tolist() choices_letter_prob = { "A": probs[0], "B": probs[1], "C": probs[2], "D": probs[3], } print(choices_letter_prob.items()) return choices_letter_prob def generator_probs(subject, question, choices_list, get_correct, model, tokenizer): # Generate the letter answer choices = [f"{chr(65+i)}. {choice}" for i, choice in enumerate(choices_list)] prompt = build_generator_prompt(subject, question, choices, get_correct) probs = get_generator_answer_probs(model, tokenizer, prompt, choices_list) return probs def load_model(model_name): model = AutoModelForCausalLM.from_pretrained( model_name, device_map="cuda", torch_dtype=torch.float16 ) tokenizer = AutoTokenizer.from_pretrained(model_name) return model, tokenizer def get_initial_generator_probs(row, model, tokenizer): gen_init = {"correct": {}, "incorrect": {}} x, y_list, subject = row["question"], row["choices"], row["subject"] for v in [True, False]: choices_letter_prob = generator_probs(subject, x, y_list, v, model, tokenizer) if v: for key, val in choices_letter_prob.items(): gen_init["correct"][key] = val else: for key, val in choices_letter_prob.items(): gen_init["incorrect"][key] = val return gen_init This is a example of my prompt > <|begin_of_text|><|start_header_id|>system<|end_header_id|> You are a helpful assistant. Answer the following multiple-choice question about college chemistry. Provide only the letter corresponding to the answer.<|eot_id|> > <|start_header_id|>user<|end_header_id|> Question: The rate, r, of a zero-order chemical reaction A → B can be expressed as which of the following? Options: A. r = k ln[A] B. r = k [A]^2 C. r = k [A] D. r = k Answer:<eot_id> > <|start_header_id|>assistant<|end_header_id|> I try to reproduce the code to test MMLU with Llama3-8B from hugging face, but get very bad performance. Continue reading...
