from peft import LoraConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch


base_model_path = "meta-llama/Llama-2-70b-hf"
model = AutoModelForCausalLM.from_pretrained(
            base_model_path,
            torch_dtype=torch.bfloat16,
            device_map="auto",
        )
peft_model_id = "haritzpuerto/LLaMA2-70B-dcot"
model.load_adapter(peft_model_id)

tokenizer = AutoTokenizer.from_pretrained(base_model_path)

[Answer 1]CoT_1
[Answer 2]CoT_2
...
[Final answer] answer

prompt = "[Question] Juan and LaKeisha roll a few objects down a ramp. They want to see which object rolls the farthest. What should they do so they can repeat their investigation?\n[Options] A) Put the objects in groups. B) Change the height of the ramp. C) Choose different objects to roll. D) Record the details of the investigation.\n[Number of answers] 2\n[Answer 1] "
inputs = tokenizer(prompt, return_tensors="pt")
output = model.generate(**inputs.to("cuda"), max_length=1024)
print(tokenizer.decode(output[0]))

<s> [Question] Juan and LaKeisha roll a few objects down a ramp. They want to see which object rolls the farthest. What should they do so they can repeat their investigation?
[Options] A) Put the objects in groups. B) Change the height of the ramp. C) Choose different objects to roll. D) Record the details of the investigation.
[Number of answers] 2
[Answer 1] 1. Juan and LaKeisha want to see which object rolls the farthest.
2. They have already rolled a few objects down the ramp.
3. To repeat their investigation, they need to do something that will affect the outcome of the experiment.
4. Putting the objects in groups will not affect the outcome of the experiment.
5. Changing the height of the ramp may affect the outcome, but it is not the best option as it requires changing the setup of the experiment.
6. Choosing different objects to roll may also affect the outcome, but it is not the best option as it does not address the issue of repeating the experiment.
7. The best option is to record the details of the investigation. This includes the objects used, the height of the ramp, and any other relevant information. By recording the details, Juan and LaKeisha can repeat the experiment with the same conditions and compare the results.
[Answer 2] Step 1: Identify the problem and the question.

Problem: Juan and LaKeisha want to see which object rolls the farthest.

Question: What should they do to repeat their investigation?

Step 2: Evaluate the options.

A) Put the objects in groups. - This option does not directly relate to the question of which object rolls the farthest, so it can be eliminated.

B) Change the height of the ramp. - This option also does not directly relate to the question of which object rolls the farthest, so it can be eliminated.

C) Choose different objects to roll. - This option is a possible solution to the question, but it does not guarantee that the object will roll the farthest.

D) Record the details of the investigation. - This option is a necessary step to repeat the investigation.

Step 3: Choose the best option.

The best option to repeat the investigation is to record the details of the investigation. This will allow them to replicate the conditions of the original experiment and compare the results.

[Final answer] D) Record the details of the investigation.</s>

@misc{puerto2024dcot,
      title={Fine-Tuning with Divergent Chains of Thought Boosts Reasoning Through Self-Correction in Language Models}, 
      author={Haritz Puerto and Tilek Chubakov and Xiaodan Zhu and Harish Tayyar Madabushi and Iryna Gurevych},
      year={2024},
      eprint={2407.03181},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2407.03181}, 
}