1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# All rights reserved.
3
4import torch
5import torch.optim as optim
6from transformers import AutoModelForCausalLM, AutoTokenizer
7from coconut import Coconut
8
9def main():
10 load_model_path = "save_models/gsm-coconut/checkpoint_22"
11 model_id = "openai-community/gpt2"
12 # load the configuration file
13 print(f"Loading from {load_model_path}.")
14
15 model = AutoModelForCausalLM.from_pretrained(model_id)
16 tokenizer = AutoTokenizer.from_pretrained(model_id)
17 tokenizer.pad_token = tokenizer.eos_token
18 tokenizer.add_tokens("<|start-latent|>")
19 tokenizer.add_tokens("<|end-latent|>")
20 tokenizer.add_tokens("<|latent|>")
21 latent_id = tokenizer.convert_tokens_to_ids("<|latent|>")
22 start_id = tokenizer.convert_tokens_to_ids("<|start-latent|>")
23 end_id = tokenizer.convert_tokens_to_ids("<|end-latent|>")
24
25 saved_weights = torch.load(
26 load_model_path, map_location=torch.device("cuda")
27 )
28
29 model.resize_token_embeddings(len(tokenizer))
30 embeddings = model.get_input_embeddings()
31 target_id = tokenizer.convert_tokens_to_ids("<<")
32 # initialize the new token embeddings with a known token
33 # it helps stablize the training
34 for token_id in [latent_id, start_id, end_id]:
35 target_embedding = embeddings.weight.data[token_id]
36 embeddings.weight.data[token_id] = target_embedding
37 # The input embeddings and lm heads are tied in GPT2. So the code below is not necessary
38 lm_head = model.lm_head
39 lm_head.weight.data[token_id] = lm_head.weight.data[target_id]
40
41 model = Coconut(model, latent_id, start_id, end_id, tokenizer.eos_token_id)
42 print(model.load_state_dict(saved_weights, strict=False))
43 model = model.to("cuda")
44
45 prompt = "Sally received the following scores on her math quizzes: 50, 80, 80. Find her mean score."
46 prompt = tokenizer(prompt, return_tensors="pt").to("cuda")
47 output = model.generate(
48 **prompt,
49 max_new_tokens=20
50 )
51 for i, o in enumerate(output):
52 print(f"Output {i}: {tokenizer.decode(o, skip_special_tokens=True)}")
53
54if __name__ == "__main__":
55 main()
56
