├── .idea
    ├── vcs.xml
    ├── .gitignore
    ├── inspectionProfiles
    │   ├── profiles_settings.xml
    │   └── Project_Default.xml
    ├── modules.xml
    └── RECAP.iml
├── requirements.txt
├── LICENSE
├── src
    ├── utils
    │   ├── retriever_utils.py
    │   └── generator_utils.py
    ├── preprocess
    │   ├── encode_comments.py
    │   ├── retrieval.py
    │   ├── recent.py
    │   └── retrieved.py
    ├── retrieve.py
    ├── models
    │   ├── retriever.py
    │   └── dialogpt.py
    ├── generate.py
    ├── train_generator.py
    ├── train_retriever.py
    └── eval.py
└── README.md


/.idea/vcs.xml:
--------------------------------------------------------------------------------
1 | <?xml version="1.0" encoding="UTF-8"?>
2 | <project version="4">
3 |   <component name="VcsDirectoryMappings">
4 |     <mapping directory="" vcs="Git" />
5 |   </component>
6 | </project>


--------------------------------------------------------------------------------
/.idea/.gitignore:
--------------------------------------------------------------------------------
1 | # Default ignored files
2 | /shelf/
3 | /workspace.xml
4 | # Editor-based HTTP Client requests
5 | /httpRequests/
6 | # Datasource local storage ignored files
7 | /dataSources/
8 | /dataSources.local.xml
9 | 


--------------------------------------------------------------------------------
/.idea/inspectionProfiles/profiles_settings.xml:
--------------------------------------------------------------------------------
1 | <component name="InspectionProjectProfileManager">
2 |   <settings>
3 |     <option name="USE_PROJECT_PROFILE" value="false" />
4 |     <version value="1.0" />
5 |   </settings>
6 | </component>


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
 1 | tqdm==4.64.1
 2 | numpy==1.23.3
 3 | nltk==3.8.1
 4 | torch==1.12.1
 5 | datasets==1.18.3
 6 | transformers==4.17.0
 7 | sentence-transformers==2.2.2
 8 | pytorch-ignite==0.4.12
 9 | bert-score==0.3.13
10 | git+https://github.com/google-research/bleurt.git


--------------------------------------------------------------------------------
/.idea/modules.xml:
--------------------------------------------------------------------------------
1 | <?xml version="1.0" encoding="UTF-8"?>
2 | <project version="4">
3 |   <component name="ProjectModuleManager">
4 |     <modules>
5 |       <module fileurl="file://$PROJECT_DIR$/.idea/RECAP.iml" filepath="$PROJECT_DIR$/.idea/RECAP.iml" />
6 |     </modules>
7 |   </component>
8 | </project>


--------------------------------------------------------------------------------
/.idea/RECAP.iml:
--------------------------------------------------------------------------------
 1 | <?xml version="1.0" encoding="UTF-8"?>
 2 | <module type="PYTHON_MODULE" version="4">
 3 |   <component name="NewModuleRootManager">
 4 |     <content url="file://$MODULE_DIR$">
 5 |       <excludeFolder url="file://$MODULE_DIR$/venv" />
 6 |     </content>
 7 |     <orderEntry type="inheritedJdk" />
 8 |     <orderEntry type="sourceFolder" forTests="false" />
 9 |   </component>
10 |   <component name="PyDocumentationSettings">
11 |     <option name="format" value="PLAIN" />
12 |     <option name="myDocStringFormat" value="Plain" />
13 |   </component>
14 | </module>


--------------------------------------------------------------------------------
/.idea/inspectionProfiles/Project_Default.xml:
--------------------------------------------------------------------------------
 1 | <component name="InspectionProjectProfileManager">
 2 |   <profile version="1.0">
 3 |     <option name="myName" value="Project Default" />
 4 |     <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
 5 |       <option name="ignoredPackages">
 6 |         <value>
 7 |           <list size="3">
 8 |             <item index="0" class="java.lang.String" itemvalue="insightface" />
 9 |             <item index="1" class="java.lang.String" itemvalue="dlib" />
10 |             <item index="2" class="java.lang.String" itemvalue="opencv_python" />
11 |           </list>
12 |         </value>
13 |       </option>
14 |     </inspection_tool>
15 |   </profile>
16 | </component>


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2023 shuailiu6626
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/src/utils/retriever_utils.py:
--------------------------------------------------------------------------------
 1 | from dataclasses import dataclass
 2 | 
 3 | import torch
 4 | from transformers import PreTrainedTokenizerBase
 5 | from typing import Any, Callable, Dict, List, NewType, Optional, Tuple, Union
 6 | from transformers.file_utils import PaddingStrategy
 7 | 
 8 | 
 9 | def encode_example(example, tokenizer):
10 |     max_len=128
11 | 
12 |     encoded_example = {}
13 |     encoded_srcs = tokenizer(
14 |         [txt.replace("<|TITLE|> ", "").replace(" <|EOS|> ", tokenizer.eos_token) for txt in example['srcs']],
15 |         max_length=max_len//2,
16 |         truncation=True,
17 |         padding="max_length",
18 |     )
19 |     encoded_tgts = tokenizer(
20 |         [txt for txt in example['tgts']],
21 |         max_length=max_len//2, 
22 |         truncation=True, 
23 |         padding="max_length"
24 |     )
25 | 
26 |     encoded_example = {
27 |         "srcs_ids": encoded_srcs.input_ids,
28 |         "srcs_attention_mask": encoded_srcs.attention_mask,
29 |         "tgts_ids": encoded_tgts.input_ids,
30 |         "tgts_attention_mask": encoded_tgts.attention_mask,
31 |     }
32 | 
33 |     return encoded_example
34 | 
35 | 
36 | @dataclass
37 | class DataCollator:
38 | 
39 |     def __call__(self, features, return_tensors=None):
40 |         import numpy as np
41 | 
42 |         seq_num = len(features[0]['srcs_ids'])
43 |         features = [{k if 'rep' not in k else 'labels': 
44 |                      torch.vstack(v) if 'rep' not in k else v.reshape(seq_num, -1) 
45 |                      for k, v in feature.items()} for feature in features]
46 |         batch = {}
47 |         for key in features[0].keys():
48 |             batch[key] = torch.cat([feature[key].unsqueeze(0) for feature in features], dim=0)
49 | 
50 |         return batch


--------------------------------------------------------------------------------
/src/preprocess/encode_comments.py:
--------------------------------------------------------------------------------
 1 | import os
 2 | import json
 3 | import argparse
 4 | from tqdm import tqdm
 5 | 
 6 | import torch
 7 | from sentence_transformers import SentenceTransformer
 8 | 
 9 | 
10 | if __name__ == '__main__':
11 | 
12 |     parser = argparse.ArgumentParser()
13 |     parser.add_argument('-d', '--data_path', type=str, help="path to the raw data", default="data/raw")
14 |     parser.add_argument('-o', '--output_path', type=str, help="output processed data path", default="temp/reps")
15 |     args = parser.parse_args()
16 | 
17 |     models = {}
18 |     models['style'] = SentenceTransformer('AnnaWegmann/Style-Embedding', device='cuda').half()
19 |     models['semantic'] = SentenceTransformer('sentence-transformers/all-mpnet-base-v2', device='cuda').half()
20 | 
21 |     for split in ["train", "valid", "test", "pandora"]:
22 |         lines = []
23 |         with open(os.path.join(args.data_path, f"{split}.jsonl")) as f:
24 |             for i, line in tqdm(enumerate(f), disable=True):
25 |                 line = json.loads(line)
26 |                 lines.append(line)
27 |         lines = sorted(lines, key=lambda x:x[1])
28 |         srcs = [line[2].replace("<|TITLE|> ", "").replace(" <|EOS|> ", "\n") for line in lines]
29 |         tgts = [line[3] for line in lines]
30 | 
31 |         for rep_type in ['style', 'semantic']:
32 |             reps = {}
33 |             reps['src'] = models[rep_type].encode(srcs, batch_size=1024, convert_to_tensor=False, normalize_embeddings=False, show_progress_bar=False)
34 |             reps['tgt'] = models[rep_type].encode(tgts, batch_size=1024, convert_to_tensor=False, normalize_embeddings=False, show_progress_bar=False)
35 |             reps['src'] = torch.from_numpy(reps['src']).half()
36 |             reps['tgt'] = torch.from_numpy(reps['tgt']).half()
37 |             torch.save(reps, os.path.join(args.output_path, f"{rep_type}/{split}_fp16.pt"), create_dir=True)


--------------------------------------------------------------------------------
/src/preprocess/retrieval.py:
--------------------------------------------------------------------------------
 1 | import os
 2 | import argparse
 3 | from tqdm import tqdm
 4 | from collections import defaultdict
 5 | import pandas as pd
 6 | import random
 7 | import json
 8 | 
 9 | import torch
10 | from datasets import load_dataset, load_from_disk
11 | 
12 | 
13 | if __name__ == '__main__':
14 | 
15 |     parser = argparse.ArgumentParser()
16 |     parser.add_argument('-d', '--data_path', type=str, help="path to the raw data", default="data/raw")
17 |     parser.add_argument('-o', '--output_path', type=str, help="output processed data path", default="data/retrieval")
18 |     args = parser.parse_args()
19 | 
20 |     for split in ["train", "valid", "test", "pandora"]:
21 |         lines = []
22 |         user2lines = defaultdict(list)
23 |         with open(os.path.join(args.data_path, f"{split}.jsonl")) as f:
24 |             for i, line in tqdm(enumerate(f), disable=True):
25 |                 line = json.loads(line)
26 |                 lines.append(line)
27 | 
28 |         lines = sorted(lines, key=lambda x:x[1])
29 |         for i, line in enumerate(lines):
30 |             author = line[0]
31 |             user2lines[author].append(i)  
32 | 
33 |         authors = sorted(user2lines.keys())
34 |         author2idx = {author: idx for idx, author in enumerate(authors)}
35 | 
36 |         author_comments = defaultdict(list)
37 |         for author in tqdm(authors, disable=True):
38 |             author_lines = user2lines.get(author, None)
39 |             assert author_lines
40 |             prev_t = 0
41 |             for author_line in author_lines:
42 |                 line = lines[author_line]
43 |                 _author, timestamp, src, tgt, subreddit = line
44 | 
45 |                 assert author == _author
46 |                 assert prev_t <= timestamp
47 |                 prev_t = timestamp
48 | 
49 |                 author_comments[author].append((int(timestamp), src, tgt, author_line))
50 |                 author_comments[author] = sorted(author_comments[author], key=lambda x:x[0])
51 |         assert len(author_comments) == len(authors)
52 | 
53 |         with open(os.path.join(args.output_path, f"{split}.jsonl"), 'w') as f:
54 |             for author, samples in tqdm(author_comments.items(), disable=True):
55 | 
56 |                 ex = {
57 |                     "srcs": [sample[1] for sample in samples],
58 |                     "tgts": [sample[2] for sample in samples],
59 |                 }
60 | 
61 |                 json_line = json.dumps(ex)
62 |                 f.write(f"{json_line}\n")


--------------------------------------------------------------------------------
/src/preprocess/recent.py:
--------------------------------------------------------------------------------
 1 | import os
 2 | import json
 3 | import argparse
 4 | from collections import defaultdict
 5 | 
 6 | 
 7 | 
 8 | if __name__ == "__main__":
 9 | 
10 |     parser = argparse.ArgumentParser()
11 |     parser.add_argument('-d', '--data_path', type=str, help="path to the raw data", default="data/raw")
12 |     parser.add_argument('-o', '--output_path', type=str, help="output processed data path", default="data/recent")
13 |     args = parser.parse_args()
14 | 
15 |     if not os.path.exists(args.output_path):
16 |         os.makedirs(args.output_path)
17 | 
18 |     for split in ["train", "valid", "test", "pandora"]:
19 |         lines = []
20 |         user2lines = defaultdict(list)
21 |         with open(os.path.join(args.data_path, f"{split}.jsonl")) as f:
22 |             for i, line in enumerate(f):
23 |                 line = json.loads(line)
24 |                 lines.append(line)
25 | 
26 |         lines = sorted(lines, key=lambda x:x[1])
27 |         for i, line in enumerate(lines):
28 |             author = line[0]
29 |             user2lines[author].append(i)  
30 | 
31 |         authors = sorted(user2lines.keys())
32 |         author2idx = {author: idx for idx, author in enumerate(authors)}
33 |         
34 |         author_comments = defaultdict(list)
35 |         for author in authors:
36 |             author_lines = user2lines.get(author, None)
37 |             assert author_lines
38 |             for author_line in author_lines:
39 |                 line = lines[author_line]
40 |                 _author, timestamp, src, tgt, subreddit = line
41 |                 assert author == _author
42 |                 author_comments[author].append((int(timestamp), src, tgt))
43 |                 author_comments[author] = sorted(author_comments[author], key=lambda x:x[0])
44 |         assert len(author_comments) == len(authors)
45 |         
46 |         global_i = 0
47 |         with open(os.path.join(args.output_path, f"{split}.jsonl"), 'w') as f:
48 |             for author, samples in author_comments.items():
49 |                 recent_responses = []
50 |                 for i, (timestamp, src, tgt) in enumerate(samples):
51 |                     ex = {
52 |                         "id": global_i,
53 |                         "author": author,
54 |                         "author_id": author2idx[author],
55 |                         "timestamp": timestamp,
56 |                         "src": src,
57 |                         "tgt": tgt,
58 |                     }
59 |                     ex["style_tgts"] = recent_responses
60 |                     ex["semantic_tgts"] = None
61 | 
62 |                     if i >= 90:
63 |                         global_i += 1
64 |                         json_line = json.dumps(ex)
65 |                         f.write(f"{json_line}\n")
66 | 
67 |                     recent_responses.insert(0, f"{tgt}")
68 |                     if len(recent_responses) > 10:
69 |                         del recent_responses[-1]


--------------------------------------------------------------------------------
/src/preprocess/retrieved.py:
--------------------------------------------------------------------------------
 1 | import os
 2 | import json
 3 | import argparse
 4 | from tqdm import tqdm
 5 | from collections import Counter, defaultdict
 6 | 
 7 | import torch
 8 | 
 9 | 
10 | if __name__ == '__main__':
11 | 
12 |     parser = argparse.ArgumentParser()
13 |     parser.add_argument('-d', '--data_path', type=str, help="path to the raw data", default="data/raw")
14 |     parser.add_argument('-r', '--retrieved_path', type=str, help="path to the retriever output", default="temp/retrieved")
15 |     parser.add_argument('-o', '--output_path', type=str, help="output processed data path", default="data/retrieved")
16 |     args = parser.parse_args()
17 | 
18 |     retrieved = {
19 |         "style": torch.load(os.path.join(args.retrieved_path, "style.pt")),
20 |         "semantic": torch.load(os.path.join(args.retrieved_path, "semantic.pt")),
21 |     }
22 | 
23 |     for split in ["train", "valid", "test", "pandora"]:
24 |         sorted_style = retrieved['style'][split]
25 |         sorted_semantic = retrieved['semantic'][split]
26 | 
27 |         lines = []
28 |         user2lines = defaultdict(list)
29 |         with open(os.path.join(args.data_path, f"{split}.jsonl")) as f:
30 |             for i, line in tqdm(enumerate(f), disable=True):
31 |                 line = json.loads(line)
32 |                 lines.append(line)
33 | 
34 |         lines = sorted(lines, key=lambda x:x[1])
35 |         for i, line in enumerate(lines):
36 |             author = line[0]
37 |             user2lines[author].append(i)  
38 | 
39 |         authors = sorted(user2lines.keys())
40 |         author2idx = {author: idx for idx, author in enumerate(authors)}
41 | 
42 |         author_comments = defaultdict(list)
43 |         for author in tqdm(authors, disable=True):
44 |             author_lines = user2lines.get(author, None)
45 |             assert author_lines
46 |             for author_line in author_lines:
47 |                 line = lines[author_line]
48 |                 _author, timestamp, src, tgt, subreddit = line
49 |                 assert author == _author
50 |                 author_comments[author].append((int(timestamp), src, tgt, author_line))
51 |                 author_comments[author] = sorted(author_comments[author], key=lambda x:x[0])
52 |         assert len(author_comments) == len(authors)
53 | 
54 |         global_i = 0
55 |         with open(os.path.join(args.output_path, f"{split}.jsonl"), 'w') as f:
56 |             for i, (author, samples) in enumerate(tqdm(author_comments.items(), disable=True)):
57 |                 most_similar_style = sorted_style[i*10:(i+1)*10]
58 |                 most_similar_semantic = sorted_semantic[i*10:(i+1)*10]
59 |                 for i, (timestamp, src, tgt, _) in enumerate(samples[-10:]):
60 |                     ex = {
61 |                         "id": global_i,
62 |                         "author": author,
63 |                         "author_id": author2idx[author],
64 |                         "timestamp": timestamp,
65 |                         "src": src,
66 |                         "tgt": tgt,
67 |                     }
68 |                     ex["style_tgts"] = [samples[j][2] for j in most_similar_style[i][:10]]
69 |                     ex["semantic_tgts"] = [samples[j][2] for j in most_similar_semantic[i][:10]]
70 | 
71 |                     global_i += 1
72 |                     json_line = json.dumps(ex)
73 |                     f.write(f"{json_line}\n")
74 | 


--------------------------------------------------------------------------------
/src/retrieve.py:
--------------------------------------------------------------------------------
 1 | import os
 2 | import argparse
 3 | import numpy as np
 4 | from tqdm import trange
 5 | from dataclasses import dataclass
 6 | from collections import defaultdict
 7 | 
 8 | import torch
 9 | from torch import nn
10 | import torch.nn.functional as F
11 | from datasets import load_from_disk
12 | from models.retriever import Retriever
13 | 
14 | 
15 | @dataclass
16 | class DataCollator:
17 | 
18 |     def __call__(self, features, return_tensors=None):
19 |         import numpy as np
20 | 
21 |         seq_num = len(features[0]['srcs_ids'])
22 |         features = [{k if 'rep' not in k else 'labels': 
23 |                      torch.vstack(v) if 'rep' not in k else v.reshape(seq_num, -1) 
24 |                      for k, v in feature.items()} for feature in features]
25 |         batch = {}
26 |         for key in features[0].keys():
27 |             batch[key] = torch.cat([feature[key].unsqueeze(0) for feature in features], dim=0)
28 | 
29 |         return batch
30 | 
31 | if __name__ == "__main__":
32 | 
33 |     parser = argparse.ArgumentParser()
34 |     parser.add_argument('-d', '--data_path', type=str, help="path to the dataset", default="data/retrieval")
35 |     parser.add_argument('-m', '--model_path', type=str, help="path to the retrieval model", default="temp/retriever")
36 |     parser.add_argument('-s', '--save_path', type=str, help="path to save the output", default="temp/retrieved")
37 |     parser.add_argument(
38 |         '-r', '--ref_type', 
39 |         type=str, default="style", 
40 |         choices=["style", "semantic"],
41 |     )
42 |     parser.add_argument(
43 |         '-b', '--batch_size', 
44 |         type=int, default=2,
45 |     )
46 |     args = parser.parse_args()
47 | 
48 |     dataset = load_from_disk(os.path.join(args.data_path, "hf_dataset"))
49 |     data_columns = [
50 |         'srcs_ids', 
51 |         'srcs_attention_mask',
52 |         'tgts_ids', 
53 |         'tgts_attention_mask',
54 |         f'{args.ref_type}_rep',
55 |     ]
56 |     dataset.set_format(
57 |         type='torch', 
58 |         columns=data_columns,
59 |     )
60 | 
61 |     data_collator = DataCollator()
62 | 
63 |     model = Retriever("distilroberta-base", use_gold_tgt_rep=False, nhead=12)
64 | 
65 |     path = os.path.join(args.model_path, args.ref_type, "model/checkpoint-best/pytorch_model.bin")
66 |     state_dict = torch.load(path)
67 |     model.load_state_dict(state_dict)
68 | 
69 |     _ = model.cuda()
70 |     _ = model.eval()
71 | 
72 |     bsz = args.batch_size
73 |     retrieved_ids = defaultdict(list)
74 |     for split in ['train', 'valid', 'test']:
75 | 
76 |         pbar = trange(0, len(dataset[split]), bsz, disable=False)
77 |         for i in pbar:
78 |             features = [dataset[split][j] for j in range(i, min(i+bsz, len(dataset[split])))]
79 |             features = data_collator(features)
80 |             batch = {k: v.cuda() for k, v in features.items()}
81 | 
82 |             preds = model(**batch).logits
83 |             preds = F.normalize(preds, dim=-1)
84 |             reps = F.normalize(batch['labels'], dim=-1)
85 | 
86 |             pred_sims = torch.matmul(preds, reps.transpose(1, 2))
87 |             pred_sims += -2. * torch.ones_like(pred_sims).triu(diagonal=90)
88 |             pred_sorted_sims, pred_sorted_indices = pred_sims.sort(descending=True, dim=-1)
89 |             retrieved_ids[split].append(pred_sorted_indices[:, :, :10].cpu())
90 | 
91 |         retrieved_ids[split] = torch.cat(retrieved_ids[split], dim=0).reshape(1, -1, 10).squeeze(0)
92 |     
93 |     torch.save(retrieved_ids, os.path.join(args.save_path, f"{args.ref_type}.pt"))


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | # RECAP
  2 | 
  3 | This is the project that implements personal specific text generation project.
  4 | It uses the gpt model as a basis and fine-tuning them to achieve the goal.
  5 | Enjoy using this proj.
  6 | # Installation
  7 | 
  8 | Commends for enviroment setup with conda.
  9 | ```bash
 10 | conda create --name recap python=3.8
 11 | conda activate recap
 12 | pip install -U pip
 13 | pip install -r requirements.txt
 14 | 
 15 | ```
 16 | 
 17 | # Data
 18 | 
 19 | The data is extracted from the Reddit dump from [pushshift.io](https://pushshift.io/). To preserve persona and personal writing style as much as possible, we did not filter out conversations with unethical content. You can download the raw data from the link [here](https://drive.google.com/file/d/1YC43Pqn15E7IIb90hjtauqRbwCOqAi3x/view?usp=sharing).
 20 | 
 21 | # Pre-processing
 22 | 
 23 | Pre-process the raw data into the format for retrieval and generation.
 24 | 
 25 | ## Retrieval Data
 26 | 
 27 | ### Encode text representations
 28 | 
 29 | ```bash
 30 | python src/preprocess/encode_comments.py -d <raw_data_path> -o <output_path>
 31 | ```
 32 | 
 33 | ### Retrieval
 34 | 
 35 | ```bash
 36 | python src/preprocess/retrieval.py -d <raw_data_path> -o <output_path>
 37 | ```
 38 | 
 39 | ## Generation Data
 40 | 
 41 | ### Most recent hisotry responses
 42 | 
 43 | ```bash
 44 | python src/preprocess/recent.py -d <raw_data_path> -o <output_path>
 45 | ```
 46 | 
 47 | ### Retrieved by hierarchical transformer
 48 | 
 49 | This requires the retriever output in `retrieved_path`. Please see section `training retriever` and `inference retrieve` for details on how to train and retrieve with the hierarchical transformer retriever.
 50 | ```bash
 51 | python src/preprocess/retrieved.py -d <raw_data_path> -r <retrieved_path> -o <output_path>
 52 | ```
 53 | 
 54 | # Training 
 55 | 
 56 | Train the retriever and the generator on a single GPU. The code works for multi GPUs, but the `batch_size` here is per device batch size, so please change it accordingly if you use more than one GPU.
 57 | 
 58 | ## Retriever
 59 | 
 60 | ```bash
 61 | python src/train_retriever.py \
 62 |     --data_path <data_path> \
 63 |     --raw_data_path <raw_data_path> \
 64 |     --reps_path <representations_path> \
 65 |     --save_path <save_path> \
 66 |     --ref_type <style OR semantic> \
 67 |     --lr 5e-5 \
 68 |     --batch_size 4 \
 69 |     --grad_accumulation 8 \
 70 |     --warmup 6250 \
 71 |     --nhead 12
 72 | ```
 73 | 
 74 | ## Generator
 75 | 
 76 | ```bash
 77 | python src/train_generator.py \
 78 |     --data_path <data_path> \
 79 |     --save_path <save_path> \
 80 |     --injection_mode <(optional) concat OR context-prefix> \
 81 |     --ref_type <(optional) style OR semantic> \
 82 |     --lr 5e-5 \
 83 |     --batch_size 128 \
 84 |     --warmup 10000
 85 | ```
 86 | 
 87 | # Inference
 88 | 
 89 | Retrieve and generate with trained models.
 90 | 
 91 | ## Retrieve
 92 | 
 93 | ```bash
 94 | python src/retrieve.py \
 95 |     --data_path <data_path> \
 96 |     --model_path <retriever_model_path> \
 97 |     --save_path <save_path> \
 98 |     --ref_type <style OR semantic>
 99 | ```
100 | 
101 | ## Generate
102 | 
103 | ```bash
104 | python src/generated.py \
105 |     --data_path <data_path> \
106 |     --model_path <generator_model_path> \
107 |     --save_path <save_path> \
108 |     --injection_mode <(optional) concat OR context-prefix> \
109 |     --ref_type <(optional) style OR semantic>
110 | ```
111 | 
112 | ## Evaluate
113 | 
114 | Please download the bleurt checkpoint BLEURT-20-D3 from [here](https://github.com/google-research/bleurt) before running the evaluation.
115 | 
116 | ```bash
117 | python src/eval.py \
118 |     --generated_path <generated_responses_path> \
119 |     --dataset_path <data_path> \
120 |     --cache_dir <eval_cache_dir> \
121 |     --cav_samples <eval_cav_samples_file>
122 | ```
123 | 


--------------------------------------------------------------------------------
/src/models/retriever.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | from torch import nn
 3 | import torch.nn.functional as F
 4 | 
 5 | from transformers import AutoTokenizer, AutoModel
 6 | from transformers.modeling_outputs import SequenceClassifierOutput
 7 | 
 8 | 
 9 | def generate_square_subsequent_mask(sz, device):
10 |     mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
11 |     mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
12 |     return mask
13 | 
14 | 
15 | class PositionalTypeEmbedding(nn.Module):
16 |     
17 |     def __init__(self, d_model=768, dropout=0.1, max_len=512):
18 |         super().__init__()
19 |         embed_dim = d_model
20 |         self.dropout = nn.Dropout(p=dropout)
21 |         self.positional_embedding = nn.Parameter(torch.Tensor(max_len, embed_dim))
22 |         nn.init.normal_(self.positional_embedding, mean=0, std=embed_dim ** -0.5)
23 |         self.type_embedding = nn.Parameter(torch.Tensor(2, embed_dim))
24 |         nn.init.normal_(self.type_embedding, mean=0, std=embed_dim ** -0.5)
25 |     
26 |     def forward(self, x):
27 |         bsz, seq_len, _ = x.size()
28 |         assert seq_len % 2 == 0
29 |         x += self.positional_embedding[:seq_len//2].repeat_interleave(2, dim=0)
30 |         x += self.type_embedding.repeat(seq_len//2, 1)
31 |         
32 |         return self.dropout(x)
33 | 
34 | 
35 | class Retriever(nn.Module):
36 |     def __init__(self, model_name, nhead=12, use_gold_tgt_rep=False):
37 |         super(Retriever, self).__init__()
38 | 
39 |         encoder_layer = nn.TransformerEncoderLayer(d_model=768, nhead=nhead, batch_first=True)
40 |         layer_norm = nn.LayerNorm(768)
41 |         self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6, norm=layer_norm).cuda()
42 |         self.positional_type_embedding = PositionalTypeEmbedding().cuda()
43 | 
44 |         self.src_encoder = AutoModel.from_pretrained(model_name)
45 |         self.tgt_encoder = None if use_gold_tgt_rep else AutoModel.from_pretrained(model_name)
46 |         self.tgt_proj = None if not use_gold_tgt_rep else nn.Linear(768, 768)
47 |         self.out_proj = nn.Linear(768, 768)
48 | 
49 |         for layer in self.transformer_encoder.layers:
50 |             for name, child in layer.named_children():
51 |                 if name == "self_attn":
52 |                     child._reset_parameters()
53 |                 elif "dropout" not in name:
54 |                     child.reset_parameters()
55 |     
56 |     def forward(
57 |         self,
58 |         srcs_ids,
59 |         srcs_attention_mask,
60 |         tgts_ids,
61 |         tgts_attention_mask,
62 |         labels,
63 |     ):
64 |         bsz, seq_num, seq_len = srcs_ids.size()
65 | 
66 |         src_reps = self.src_encoder(
67 |             input_ids=srcs_ids.reshape(bsz*seq_num, -1), 
68 |             attention_mask=srcs_attention_mask.reshape(bsz*seq_num, -1)
69 |         ).last_hidden_state.mean(1).reshape(bsz, seq_num, -1)
70 | 
71 |         if self.tgt_encoder is None:
72 |             tgt_reps = self.tgt_proj(labels)
73 |         else:
74 |             tgt_reps = self.tgt_encoder(
75 |                 input_ids=tgts_ids.reshape(bsz*seq_num, -1), 
76 |                 attention_mask=tgts_attention_mask.reshape(bsz*seq_num, -1)
77 |             ).last_hidden_state.mean(1).reshape(bsz, seq_num, -1)
78 | 
79 |         transformer_input = torch.cat([src_reps, tgt_reps], dim=-1).reshape(bsz, seq_num*2, -1)
80 |         transformer_input = self.positional_type_embedding(transformer_input)
81 |         causal_mask = generate_square_subsequent_mask(seq_num*2, transformer_input.device)
82 |         transformer_output = self.transformer_encoder(transformer_input, mask=causal_mask)
83 |         transformer_output = self.out_proj(transformer_output[: ,::2, :])
84 | 
85 |         transformer_output = transformer_output[:, -10:, :]
86 |         labels = labels[:, -10:, :]
87 | 
88 |         loss_fct = nn.CosineEmbeddingLoss()
89 |         loss = loss_fct(
90 |             transformer_output.reshape(bsz*seq_num//10, -1), 
91 |             labels.reshape(bsz*seq_num//10, -1), 
92 |             target=torch.ones(bsz*seq_num//10).to(transformer_output.device)
93 |         )
94 | 
95 |         return SequenceClassifierOutput(
96 |             loss=loss,
97 |             logits=transformer_output,
98 |             hidden_states=None,
99 |         )


--------------------------------------------------------------------------------
/src/generate.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import argparse
  3 | from collections import defaultdict
  4 | import json
  5 | from tqdm.auto import tqdm, trange
  6 | 
  7 | import torch
  8 | 
  9 | from transformers import (
 10 |     AutoTokenizer,
 11 |     AutoModelForCausalLM,
 12 |     Trainer,
 13 |     TrainingArguments,
 14 |     EarlyStoppingCallback,
 15 |     GPT2LMHeadModel,
 16 |     set_seed,
 17 | )
 18 | 
 19 | from datasets import load_from_disk
 20 | 
 21 | from dataclasses import dataclass
 22 | from transformers import PreTrainedTokenizerBase
 23 | from typing import Any, Callable, Dict, List, NewType, Optional, Tuple, Union
 24 | from transformers.file_utils import PaddingStrategy
 25 | 
 26 | from utils.generator_utils import encode_example, DataCollator, generate_responses
 27 | 
 28 | 
 29 | if __name__ == "__main__":
 30 | 
 31 |     parser = argparse.ArgumentParser()
 32 |     parser.add_argument('-d', '--data_path', type=str, help="path to the dataset", default="data/recent")
 33 |     parser.add_argument('-m', '--model_path', type=str, help="path to trained model", default="temp/model")
 34 |     parser.add_argument('-s', '--save_path', type=str, help="path to save generated texts", default="temp/generated")
 35 |     parser.add_argument(
 36 |         '-i', '--injection_mode', 
 37 |         type=str, default=None, 
 38 |         choices=["concat", "prefix", "context-prefix"],
 39 |     )
 40 |     parser.add_argument(
 41 |         '-r', '--ref_type', 
 42 |         type=str, default=None, 
 43 |         choices=["style", "semantic", "mixed"],
 44 |     )
 45 |     parser.add_argument(
 46 |         '-n', '--num_ref', 
 47 |         type=int, default=10,
 48 |     )
 49 |     parser.add_argument(
 50 |         '-b', '--batch_size', 
 51 |         type=int, default=128,
 52 |     )
 53 |     parser.add_argument('--tqdm', action='store_true')
 54 |     args = parser.parse_args()
 55 | 
 56 |     task = f"{args.injection_mode}-{args.ref_type}"
 57 | 
 58 |     tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small", use_fast=True)
 59 |     tokenizer.pad_token = tokenizer.eos_token
 60 |     roberta_tokenizer = AutoTokenizer.from_pretrained("roberta-base", use_fast=True)
 61 | 
 62 |     data_dir = args.data_path
 63 |     hf_dataset_path = os.path.join(data_dir, "hf_dataset")
 64 |     checkpoint_path = os.path.join(args.model_path, "checkpoint-best")
 65 |     
 66 |     # load hf dataset
 67 |     data_columns = [
 68 |         'id',
 69 |         'author_id',
 70 |         'input_ids', 
 71 |         'attention_mask', 
 72 |         'labels', 
 73 |     ]
 74 | 
 75 |     if args.injection_mode == "concat":
 76 |         assert args.ref_type != 'mixed', "ref_type == mixed only supports prefix based models"
 77 |         from models.dialogpt import GPT2LMHeadModel as PRG
 78 |         data_columns += [
 79 |             f'decoder_{args.ref_type}_ids', 
 80 |             f'decoder_{args.ref_type}_attention_mask',
 81 |         ]
 82 |     elif args.ref_type is not None:
 83 |         if args.injection_mode == "prefix":
 84 |             from models.dialogpt import PrefixGPT2 as PRG
 85 |         elif args.injection_mode == "context-prefix":
 86 |             from models.dialogpt import ContextPrefixGPT2 as PRG
 87 |             data_columns += ['context_ids', 'context_attention_mask']
 88 |         
 89 |         if args.ref_type != 'mixed':
 90 |             data_columns += [
 91 |                 f'{args.ref_type}_ids', 
 92 |                 f'{args.ref_type}_attention_mask',
 93 |             ]
 94 |         else:
 95 |             data_columns += [
 96 |                 'style_ids', 
 97 |                 'style_attention_mask',
 98 |                 'semantic_ids', 
 99 |                 'semantic_attention_mask',
100 |             ]
101 | 
102 |     dataset = load_from_disk(f"{data_dir}/hf_dataset")
103 |     print(f"data loaded from {data_dir}")
104 |     
105 |     dataset.set_format(
106 |         type='torch', 
107 |         columns=data_columns,
108 |     )
109 | 
110 |     if args.injection_mode is None:
111 |         model = GPT2LMHeadModel.from_pretrained(checkpoint_path, local_files_only=True)
112 |     else:
113 |         model = PRG.from_pretrained(checkpoint_path, local_files_only=True)
114 | 
115 |     max_style_num = args.num_ref
116 |     max_semantic_num = args.num_ref
117 |     data_collator = DataCollator(
118 |         tokenizer, 
119 |         roberta_tokenizer,
120 |         model,
121 |         max_style_num=max_style_num,
122 |         max_semantic_num=max_semantic_num,
123 |     )
124 |     
125 |     lr = 5e-5
126 |     training_args = TrainingArguments(
127 |         save_strategy="no",
128 |         output_dir="./temp/temp", 
129 |         overwrite_output_dir=True,
130 |         num_train_epochs=10, 
131 |         per_device_train_batch_size=args.batch_size,  
132 |         per_device_eval_batch_size=args.batch_size, 
133 |         gradient_accumulation_steps=1,
134 |         learning_rate=lr,
135 |         label_smoothing_factor=0.0,
136 |         warmup_steps=0,
137 |         weight_decay=0.01, 
138 |         fp16=False,
139 |         logging_dir=None,
140 |         evaluation_strategy="steps",
141 |         logging_first_step=False,
142 |         logging_steps=2000,    # eval_steps is default to this value
143 |         save_steps=2000,
144 |         save_total_limit=1,
145 |         do_eval=True, 
146 |         do_predict=False,
147 |         metric_for_best_model="loss",
148 |         load_best_model_at_end=False,
149 |         disable_tqdm=not args.tqdm,
150 |         group_by_length=False,
151 |         length_column_name="length",
152 |         dataloader_num_workers=16,
153 |         remove_unused_columns=False,
154 |     )
155 | 
156 |     trainer = Trainer(
157 |         model=model,
158 |         args=training_args,
159 |         train_dataset=dataset["train"],
160 |         eval_dataset=dataset["valid"],
161 |         tokenizer=tokenizer,
162 |         data_collator=data_collator,
163 |         callbacks=[EarlyStoppingCallback(early_stopping_patience=10)],
164 |     )
165 | 
166 |     set_seed(42)
167 |     if not os.path.exists(args.save_path):
168 |         os.makedirs(args.save_path)
169 |     for p_prob in [0.8]:
170 |         generated = generate_responses(
171 |             trainer,
172 |             dataset["test"],
173 |             max_length=128, 
174 |             num_return_sequences=1, 
175 |             do_sample=True,
176 |             top_p=p_prob,
177 |             top_k=0,
178 |             disable_tqdm=not args.tqdm
179 |         )
180 |         
181 |         with open(os.path.join(args.save_path, "generated.jsonl"), "w") as f:
182 |             for ex in generated:
183 |                 f.write(json.dumps(ex) + '\n')
184 | 


--------------------------------------------------------------------------------
/src/train_generator.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import shutil
  3 | import argparse
  4 | 
  5 | from utils.generator_utils import encode_example, DataCollator
  6 | 
  7 | from datasets import load_dataset, load_from_disk
  8 | 
  9 | from transformers import (
 10 |     AutoModelForCausalLM,
 11 |     AutoTokenizer,
 12 |     Trainer,
 13 |     TrainingArguments,
 14 |     EarlyStoppingCallback,
 15 | )
 16 | 
 17 | if __name__ == "__main__":
 18 | 
 19 |     local_rank = int(os.environ.get("LOCAL_RANK", -1))
 20 | 
 21 |     parser = argparse.ArgumentParser()
 22 |     parser.add_argument('-d', '--data_path', type=str, help="path to the dataset", default="data/recent")
 23 |     parser.add_argument('-s', '--save_path', type=str, help="path to save trained model", default="temp/generator")
 24 |     parser.add_argument(
 25 |         '-i', '--injection_mode', 
 26 |         type=str, default=None, 
 27 |         choices=["concat", "prefix", "context-prefix"],
 28 |     )
 29 |     parser.add_argument(
 30 |         '-r', '--ref_type', 
 31 |         type=str, default=None, 
 32 |         choices=["style", "semantic", "mixed"],
 33 |     )
 34 |     parser.add_argument(
 35 |         '-n', '--num_ref', 
 36 |         type=int, default=10,
 37 |     )
 38 |     parser.add_argument(
 39 |         '-l', '--lr', 
 40 |         type=float, default=5e-5,
 41 |     )
 42 |     parser.add_argument(
 43 |         '-b', '--batch_size', 
 44 |         type=int, default=128,
 45 |     )
 46 |     parser.add_argument(
 47 |         '-g', '--grad_accumulation', 
 48 |         type=int, default=1,
 49 |     )
 50 |     parser.add_argument(
 51 |         '-w', '--warmup', 
 52 |         type=int, default=0,
 53 |     )
 54 |     parser.add_argument('--tqdm', action='store_true')
 55 |     args = parser.parse_args()
 56 | 
 57 |     task = f"{args.injection_mode}-{args.ref_type}"
 58 | 
 59 |     tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small", use_fast=True)
 60 |     tokenizer.pad_token = tokenizer.eos_token
 61 |     roberta_tokenizer = AutoTokenizer.from_pretrained("roberta-base", use_fast=True)
 62 | 
 63 |     data_dir = args.data_path
 64 |     hf_dataset_path = os.path.join(data_dir, "hf_dataset")
 65 | 
 66 |     # tokenize raw data and save if needed
 67 |     if local_rank <= 0 and not os.path.exists(hf_dataset_path):
 68 |         print(f"tokenizing data in {data_dir} ...")
 69 | 
 70 |         cache_dir = os.path.join(data_dir, "cache")
 71 | 
 72 |         if os.path.exists(cache_dir):
 73 |             shutil.rmtree(cache_dir)
 74 | 
 75 |         data_files = {
 76 |             "train": "train.jsonl",
 77 |             "valid": "valid.jsonl",
 78 |             "test": "test.jsonl",
 79 |         }
 80 | 
 81 |         dataset = load_dataset(
 82 |             data_dir, 
 83 |             data_files=data_files, 
 84 |             cache_dir=cache_dir,
 85 |         )
 86 | 
 87 |         encoded_dataset = dataset.map(encode_example, num_proc=16, fn_kwargs={"tokenizer": tokenizer, "roberta_tokenizer": roberta_tokenizer})
 88 |         encoded_dataset.save_to_disk(hf_dataset_path)
 89 |         shutil.rmtree(cache_dir)
 90 |         del dataset, encoded_dataset
 91 |     
 92 |     # load hf dataset
 93 |     data_columns = [
 94 |         'input_ids', 
 95 |         'attention_mask', 
 96 |         'labels', 
 97 |     ]
 98 | 
 99 |     if args.injection_mode == "concat":
100 |         assert args.ref_type != 'mixed', "ref_type == mixed only supports prefix based models"
101 |         from models.dialogpt import GPT2LMHeadModel as PRG
102 |         data_columns += [
103 |             f'decoder_{args.ref_type}_ids', 
104 |             f'decoder_{args.ref_type}_attention_mask',
105 |         ]
106 |     elif args.ref_type is not None:
107 |         if args.injection_mode == "prefix":
108 |             from models.dialogpt import PrefixGPT2 as PRG
109 |         elif args.injection_mode == "context-prefix":
110 |             from models.dialogpt import ContextPrefixGPT2 as PRG
111 |             data_columns += ['context_ids', 'context_attention_mask']
112 |         
113 |         if args.ref_type != 'mixed':
114 |             data_columns += [
115 |                 f'{args.ref_type}_ids', 
116 |                 f'{args.ref_type}_attention_mask',
117 |             ]
118 |         else:
119 |             data_columns += [
120 |                 'style_ids', 
121 |                 'style_attention_mask',
122 |                 'semantic_ids', 
123 |                 'semantic_attention_mask',
124 |             ]
125 | 
126 |     dataset = load_from_disk(f"{data_dir}/hf_dataset")
127 |     print(f"data loaded from {data_dir}")
128 |     
129 |     dataset.set_format(
130 |         type='torch', 
131 |         columns=data_columns,
132 |     )
133 | 
134 |     if args.injection_mode is None:
135 |         model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
136 |     else:
137 |         model = PRG.from_pretrained("microsoft/DialoGPT-small")
138 | 
139 |     max_style_num = args.num_ref
140 |     max_semantic_num = args.num_ref
141 |     data_collator = DataCollator(
142 |         tokenizer, 
143 |         roberta_tokenizer,
144 |         model,
145 |         max_style_num=max_style_num,
146 |         max_semantic_num=max_semantic_num,
147 |     )
148 | 
149 |     save_path = args.save_path
150 |     training_args = TrainingArguments(
151 |         output_dir=f"{save_path}/model", 
152 |         overwrite_output_dir=True,
153 |         num_train_epochs=10, 
154 |         per_device_train_batch_size=args.batch_size,  
155 |         per_device_eval_batch_size=args.batch_size, 
156 |         gradient_accumulation_steps=args.grad_accumulation,
157 |         learning_rate=args.lr,
158 |         label_smoothing_factor=0.0,
159 |         warmup_steps=args.warmup,
160 |         weight_decay=0.01, 
161 |         fp16=True,
162 |         logging_dir=f"{save_path}/log",
163 |         evaluation_strategy="steps",
164 |         logging_first_step=False,
165 |         logging_steps=100,    # eval_steps is default to this value
166 |         eval_steps=1000,
167 |         save_steps=1000,
168 |         save_total_limit=1,
169 |         do_eval=True, 
170 |         do_predict=False,
171 |         metric_for_best_model="loss",
172 |         load_best_model_at_end=True,
173 |         disable_tqdm=not args.tqdm,
174 |         group_by_length=False,
175 |         length_column_name="length",
176 |         dataloader_num_workers=8,
177 |         remove_unused_columns=False,
178 |     )
179 | 
180 |     trainer = Trainer(
181 |         model=model,
182 |         args=training_args,
183 |         train_dataset=dataset["train"],
184 |         eval_dataset=dataset["valid"],
185 |         tokenizer=tokenizer,
186 |         data_collator=data_collator,
187 |         callbacks=[EarlyStoppingCallback(early_stopping_patience=10)],
188 |     )
189 | 
190 | 
191 |     pytorch_total_params = sum(p.numel() for p in model.parameters())
192 |     print(f"# Params: {pytorch_total_params}")
193 | 
194 |     print(
195 |         f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
196 |         + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
197 |     )
198 | 
199 |     print(trainer.train())
200 |     print("finish training")
201 |     print("evaluate on validation set")
202 |     print(trainer.evaluate())
203 |     print("evaluate on test set")
204 |     print(trainer.evaluate(dataset["test"]))
205 | 
206 | 


--------------------------------------------------------------------------------
/src/train_retriever.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import json
  3 | import shutil
  4 | import argparse
  5 | from collections import defaultdict
  6 | from tqdm import tqdm
  7 | 
  8 | from datasets import load_dataset, load_from_disk
  9 | from tokenizers import Tokenizer
 10 | 
 11 | import torch
 12 | from torch import nn
 13 | from transformers import (
 14 |     AutoTokenizer,
 15 |     AutoModel,
 16 |     Trainer,
 17 |     TrainingArguments,
 18 |     EarlyStoppingCallback,
 19 | )
 20 | 
 21 | from dataclasses import dataclass
 22 | from transformers import PreTrainedTokenizerBase
 23 | from typing import Any, Callable, Dict, List, NewType, Optional, Tuple, Union
 24 | from transformers.file_utils import PaddingStrategy
 25 | 
 26 | from utils.retriever_utils import encode_example, DataCollator
 27 | from models.retriever import Retriever
 28 | 
 29 | 
 30 | if __name__ == "__main__":
 31 | 
 32 |     local_rank = int(os.environ.get("LOCAL_RANK", -1))
 33 | 
 34 |     parser = argparse.ArgumentParser()
 35 |     parser.add_argument('-d', '--data_path', type=str, help="path to the dataset", default="data/retrieval")
 36 |     parser.add_argument('-u', '--raw_data_path', type=str, help="path to the raw data", default="data/raw")
 37 |     parser.add_argument('-e', '--reps_path', type=str, help="path to pre-encoded representations", default="temp/reps")
 38 |     parser.add_argument('-s', '--save_path', type=str, help="path to save trained model", default="temp/retriever")
 39 |     parser.add_argument(
 40 |         '-r', '--ref_type', 
 41 |         type=str, default="style", 
 42 |         choices=["style", "semantic"],
 43 |     )
 44 |     parser.add_argument(
 45 |         '-b', '--batch_size', 
 46 |         type=int, default=8,
 47 |     )
 48 |     parser.add_argument(
 49 |         '-g', '--grad_accumulation', 
 50 |         type=int, default=1,
 51 |     )
 52 |     parser.add_argument(
 53 |         '-l', '--lr', 
 54 |         type=float, default=5e-5,
 55 |     )
 56 |     parser.add_argument(
 57 |         '-w', '--warmup', 
 58 |         type=int, default=0,
 59 |     )
 60 |     parser.add_argument('--tqdm', action='store_true')
 61 |     parser.add_argument(
 62 |         '--nhead', 
 63 |         type=int, default=12,
 64 |     )
 65 |     args = parser.parse_args()
 66 |     
 67 |     data_dir = args.data_path
 68 |     hf_dataset_path = os.path.join(data_dir, "hf_dataset")
 69 | 
 70 |     tokenizer = AutoTokenizer.from_pretrained("roberta-base", use_fast=True, local_files_only=True)
 71 | 
 72 |     # preprocess data and save if needed
 73 |     if local_rank <= 0 and not os.path.exists(hf_dataset_path):
 74 |         cache_dir = os.path.join(data_dir, "cache")
 75 |         # tokenize data
 76 |         print(f"tokenize data in {data_dir}")
 77 |         data_files = {
 78 |             "train": "train.jsonl",
 79 |             "valid": "valid.jsonl",
 80 |             "test": "test.jsonl",
 81 |         }
 82 |         dataset = load_dataset(
 83 |             data_dir, 
 84 |             data_files=data_files, 
 85 |             cache_dir=cache_dir,
 86 |         )
 87 |         encoded_dataset = dataset.map(encode_example, num_proc=16, fn_kwargs={"tokenizer": tokenizer})
 88 |         encoded_dataset.save_to_disk(f"{data_dir}/hf_dataset")
 89 |         shutil.rmtree(cache_dir)
 90 |         del dataset, encoded_dataset
 91 | 
 92 |         # add pre-encoded representation vectors
 93 |         reps = defaultdict(lambda: defaultdict(list))
 94 | 
 95 |         for split in ["train", "valid", "test"]:
 96 |             lines = []
 97 |             user2lines = defaultdict(list)
 98 |             with open(os.path.join(args.raw_data_path, f"{split}.jsonl")) as f:
 99 |                 for i, line in tqdm(enumerate(f), disable=True):
100 |                     line = json.loads(line)
101 |                     lines.append(line)
102 | 
103 |             lines = sorted(lines, key=lambda x:x[1])
104 |             for i, line in enumerate(lines):
105 |                 author = line[0]
106 |                 user2lines[author].append(i)  
107 | 
108 |             authors = sorted(user2lines.keys())
109 |             author2idx = {author: idx for idx, author in enumerate(authors)}
110 | 
111 |             author_comments = defaultdict(list)
112 |             for author in tqdm(authors, disable=True):
113 |                 author_lines = user2lines.get(author, None)
114 |                 assert author_lines
115 |                 prev_t = 0
116 |                 for author_line in author_lines:
117 |                     line = lines[author_line]
118 |                     _author, timestamp, src, tgt, subreddit = line
119 | 
120 |                     assert author == _author
121 |                     assert prev_t <= timestamp
122 |                     prev_t = timestamp
123 | 
124 |                     author_comments[author].append((int(timestamp), src, tgt, author_line))
125 |                     author_comments[author] = sorted(author_comments[author], key=lambda x:x[0])
126 |             assert len(author_comments) == len(authors)
127 | 
128 |             style_reps = torch.load(os.path.join(args.reps_path, f"style/{split}_fp16.pt"))['tgt']
129 |             semantic_reps = torch.load(os.path.join(args.reps_path, f"semantic/{split}_fp16.pt"))['tgt']
130 |             style_reps = torch.nn.functional.normalize(style_reps, dim=-1)
131 |             semantic_reps = torch.nn.functional.normalize(semantic_reps, dim=-1)
132 | 
133 |             for author, samples in tqdm(author_comments.items(), disable=True):
134 | 
135 |                 glb_ids = torch.tensor([sample[3] for sample in samples])
136 |                 tgt_style_reps = style_reps[glb_ids]
137 |                 tgt_semantic_reps = semantic_reps[glb_ids]     
138 |                 
139 |                 reps['style'][split].append(tgt_style_reps.cpu())
140 |                 reps['semantic'][split].append(tgt_semantic_reps.cpu())
141 |             
142 |         reps = {
143 |             k0: {
144 |                 k1: [
145 |                     item.flatten().float().numpy() for item in v1
146 |                 ] for k1, v1 in v0.items()
147 |             } for k0, v0 in reps.items()
148 |         }
149 | 
150 |         encoded_dataset = load_from_disk(f"{data_dir}/hf_dataset")
151 | 
152 |         for rep_type in reps.keys():
153 |             for split in reps[rep_type].keys():
154 |                 print(rep_type, split)
155 |                 encoded_dataset[split] = encoded_dataset[split].add_column(f"{rep_type}_rep", reps[rep_type][split])
156 |         
157 |         shutil.rmtree(f"{data_dir}/hf_dataset")
158 |         encoded_dataset.save_to_disk(f"{data_dir}/hf_dataset")
159 | 
160 |     data_columns = [
161 |         'srcs_ids', 
162 |         'srcs_attention_mask', 
163 |         'tgts_ids', 
164 |         'tgts_attention_mask',
165 |         f'{args.ref_type}_rep',
166 |     ]
167 |     dataset = load_from_disk(f"{data_dir}/hf_dataset")
168 |     print(f"data loaded from {data_dir}")
169 |     dataset.set_format(
170 |         type='torch', 
171 |         columns=data_columns,
172 |     )
173 |     model = Retriever("distilroberta-base", nhead=args.nhead, use_gold_tgt_rep=False)
174 | 
175 |     data_collator = DataCollator()
176 |     
177 |     lr = args.lr
178 |     save_path = args.save_path
179 |     training_args = TrainingArguments(
180 |         output_dir=f"{save_path}/{args.ref_type}/model", 
181 |         overwrite_output_dir=True,
182 |         num_train_epochs=10,
183 |         per_device_train_batch_size=args.batch_size,  
184 |         per_device_eval_batch_size=args.batch_size, 
185 |         gradient_accumulation_steps=args.grad_accumulation,
186 |         learning_rate=lr,
187 |         label_smoothing_factor=0.0,
188 |         warmup_steps=args.warmup,
189 |         weight_decay=0.01, 
190 |         fp16=True,
191 |         logging_dir=f"{save_path}/{args.ref_type}/log",
192 |         evaluation_strategy="steps",
193 |         logging_first_step=False,
194 |         logging_steps=125,    # eval_steps is default to this value
195 |         eval_steps=625,
196 |         save_steps=625,
197 |         save_total_limit=1,
198 |         do_eval=True, 
199 |         do_predict=False,
200 |         metric_for_best_model="loss",
201 |         load_best_model_at_end=True,
202 |         disable_tqdm=not args.tqdm,
203 |         group_by_length=False,
204 |         length_column_name="length",
205 |         dataloader_num_workers=8,
206 |         remove_unused_columns=False,
207 |     )
208 | 
209 |     trainer = Trainer(
210 |         model=model,
211 |         args=training_args,
212 |         train_dataset=dataset["train"],
213 |         eval_dataset=dataset["valid"],
214 |         tokenizer=tokenizer,
215 |         data_collator=data_collator,
216 |         callbacks=[EarlyStoppingCallback(early_stopping_patience=10)],
217 |     )
218 | 
219 |     pytorch_total_params = sum(p.numel() for p in model.parameters())
220 |     print(f"# Params: {pytorch_total_params}")
221 | 
222 |     print(
223 |         f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
224 |         + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
225 |     )
226 | 
227 |     print(trainer.train())
228 |     print("finish training")
229 |     print("evaluate on validation set")
230 |     print(trainer.evaluate())
231 |     print("evaluate on test set")
232 |     print(trainer.evaluate(dataset["test"]))


--------------------------------------------------------------------------------
/src/eval.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import json
  3 | import copy
  4 | import random
  5 | import argparse
  6 | import numpy as np
  7 | from tqdm.auto import trange, tqdm
  8 | from collections import defaultdict, Counter
  9 | from nltk.tokenize import word_tokenize
 10 | from nltk.translate.bleu_score import sentence_bleu, corpus_bleu
 11 | from ignite.metrics import RougeL
 12 | from nltk.translate.meteor_score import meteor_score
 13 | from bleurt import score as bleurt_score
 14 | from bert_score import score as compute_bert_score
 15 | 
 16 | import torch
 17 | from torch import nn
 18 | 
 19 | from transformers import set_seed
 20 | from datasets import load_from_disk
 21 | 
 22 | from sentence_transformers import SentenceTransformer
 23 | 
 24 | 
 25 | def distinct(seqs):
 26 |     """ Calculate intra/inter distinct 1/2. """
 27 |     batch_size = len(seqs)
 28 |     intra_dist1, intra_dist2 = [], []
 29 |     unigrams_all, bigrams_all = Counter(), Counter()
 30 |     for seq in seqs:
 31 |         unigrams = Counter(seq)
 32 |         bigrams = Counter(zip(seq, seq[1:]))
 33 |         intra_dist1.append((len(unigrams)+1e-12) / (len(seq)+1e-5))
 34 |         intra_dist2.append((len(bigrams)+1e-12) / (max(0, len(seq)-1)+1e-5))
 35 | 
 36 |         unigrams_all.update(unigrams)
 37 |         bigrams_all.update(bigrams)
 38 | 
 39 |     inter_dist1 = (len(unigrams_all)+1e-12) / (sum(unigrams_all.values())+1e-5)
 40 |     inter_dist2 = (len(bigrams_all)+1e-12) / (sum(bigrams_all.values())+1e-5)
 41 |     intra_dist1 = np.average(intra_dist1)
 42 |     intra_dist2 = np.average(intra_dist2)
 43 |     return intra_dist1, intra_dist2, inter_dist1, inter_dist2
 44 | 
 45 | 
 46 | def corpus_meteor(refs, cands):
 47 |     scores = []
 48 |     for ref, cand in zip(refs, cands):
 49 |         scores.append(meteor_score(ref, cand))
 50 |     return np.mean(scores)
 51 | 
 52 | 
 53 | if __name__ == "__main__":
 54 | 
 55 |     parser = argparse.ArgumentParser()
 56 |     parser.add_argument('-g', '--generated_path', type=str, required=True)
 57 |     parser.add_argument('-d', '--dataset_path', type=str, required=True)
 58 |     parser.add_argument('-c', '--cache_dir', type=str, required=True)
 59 |     parser.add_argument('-v', '--cav_samples', type=str, required=True)
 60 |     args = parser.parse_args()
 61 | 
 62 |     # preprocess and save if needed
 63 |     if not os.path.exists(args.cache_dir):
 64 |         print("Encoding and preprocess labels")
 65 | 
 66 |         dataset = load_from_disk(args.dataset_path + '/test')
 67 | 
 68 |         generated = []
 69 |         with open(args.generated_path) as f:
 70 |             for line in f:
 71 |                 generated.append(json.loads(line))
 72 | 
 73 |         txts = []
 74 |         for i in range(len(dataset)):
 75 |             txts.append(generated[i]["label"])
 76 | 
 77 |         assert len(generated) == len(txts)
 78 | 
 79 |         model = SentenceTransformer('AnnaWegmann/Style-Embedding', device='cuda')
 80 |         style_vecs = model.encode(txts, batch_size=512, show_progress_bar=True)
 81 |         model = SentenceTransformer('all-mpnet-base-v2', device='cuda')
 82 |         semantic_vecs = model.encode(txts, batch_size=512, show_progress_bar=True)
 83 | 
 84 |         author_txts = defaultdict(list)
 85 |         author_style_vecs = defaultdict(list)
 86 |         author_semantic_vecs = defaultdict(list)
 87 |         line2author = []
 88 |         for i in trange(len(dataset), disable=True):
 89 |             author_id = dataset[i]["author"]
 90 |             line2author.append(author_id)
 91 |             author_txts[author_id].append(txts[i])
 92 |             author_style_vecs[author_id].append(torch.tensor(style_vecs[i]))
 93 |             author_semantic_vecs[author_id].append(torch.tensor(semantic_vecs[i]))
 94 | 
 95 |         if not os.path.exists(args.cache_dir):
 96 |             os.makedirs(args.cache_dir)
 97 | 
 98 |         torch.save(author_style_vecs, os.path.join(args.cache_dir, "style_vecs.pt"))
 99 |         torch.save(author_semantic_vecs, os.path.join(args.cache_dir, "semantic_vecs.pt"))
100 |         with open(os.path.join(args.cache_dir, "txts.json"), "w") as f:
101 |             json.dump(author_txts, f)
102 |         with open(os.path.join(args.cache_dir, "line2author.json"), "w") as f:
103 |             json.dump(line2author, f)
104 | 
105 |     # evaluate
106 |     print("Evaluating")
107 |     label_style_vecs = torch.load(os.path.join(args.cache_dir, "style_vecs.pt"))
108 |     label_semantic_vecs = torch.load(os.path.join(args.cache_dir, "semantic_vecs.pt"))
109 |     with open(os.path.join(args.cache_dir, "txts.json")) as f:
110 |         label_txts = json.load(f)
111 |     with open(args.cav_samples) as f:
112 |         dc_cav_samples = json.load(f)
113 |     author_ids = sorted(label_txts.keys())
114 |     with open(os.path.join(args.cache_dir, "line2author.json")) as f:
115 |         line2author = json.load(f)
116 | 
117 |     cossim = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
118 |     set_seed(42)
119 | 
120 |     full_path = args.generated_path
121 | 
122 |     generated = []
123 |     with open(full_path) as f:
124 |         for line in f:
125 |             generated.append(json.loads(line))
126 | 
127 |     txts = []
128 |     for i in range(len(line2author)):
129 |         txts.append(generated[i]["generated"][0])
130 | 
131 |     model = SentenceTransformer('AnnaWegmann/Style-Embedding', device='cuda')
132 |     style_vecs = model.encode(txts, batch_size=128, show_progress_bar=True)
133 |     model = SentenceTransformer('all-mpnet-base-v2', device='cuda')
134 |     semantic_vecs = model.encode(txts, batch_size=128, show_progress_bar=True)
135 | 
136 |     gen_txts = defaultdict(list)
137 |     gen_style_vecs = defaultdict(list)
138 |     gen_semantic_vecs = defaultdict(list)
139 | 
140 |     for i in trange(len(line2author), disable=True):
141 |         author_id = line2author[i]
142 |         gen_txts[author_id].append(txts[i])
143 |         gen_style_vecs[author_id].append(torch.tensor(style_vecs[i]))
144 |         gen_semantic_vecs[author_id].append(torch.tensor(semantic_vecs[i]))
145 | 
146 |     combined_txts = []
147 |     combined_refs = []
148 | 
149 |     for author_id in tqdm(author_ids, disable=True):
150 | 
151 |         combined_refs += label_txts[author_id]
152 |         combined_txts += gen_txts[author_id]
153 | 
154 |     # BLEURT
155 |     scorer = bleurt_score.BleurtScorer('BLEURT-20-D3') # please download the bleurt checkpoint from https://github.com/google-research/bleurt
156 |     bleurt_result = scorer.score(references=combined_refs, candidates=combined_txts)
157 |     print(f"BELURT: {np.mean(bleurt_result):.4f}")
158 | 
159 |     # BERTScore
160 |     print("Computing BERTScore...")
161 |     (P, R, F), hashname = compute_bert_score(combined_txts, combined_refs, lang="en", return_hash=True)
162 |     print(f"BERTScore: {float(F.mean()):.4f}")
163 | 
164 |     # BLEU
165 |     combined_txts = [word_tokenize(item) for item in tqdm(combined_txts, disable=True)]
166 |     combined_refs = [[word_tokenize(item)] for item in tqdm(combined_refs, disable=True)]
167 | 
168 |     bleu_1 = corpus_bleu(combined_refs, combined_txts, weights=(1.0,))
169 |     bleu_2 = corpus_bleu(combined_refs, combined_txts, weights=(0.5, 0.5))
170 |     bleu_3 = corpus_bleu(combined_refs, combined_txts, weights=(1./3, 1./3, 1./3))
171 |     bleu_4 = corpus_bleu(combined_refs, combined_txts, weights=(0.25, 0.25, 0.25, 0.25))
172 |     print(f"BLEU-(1|2|3|4): {bleu_1:.4f} | {bleu_2:.4f} | {bleu_3:.4f} | {bleu_4:.4f}")
173 | 
174 |     # rouge-l
175 |     rouge_l = RougeL()
176 |     rouge_l.update((combined_txts, combined_refs))
177 |     rouge_l_score = rouge_l.compute()
178 |     print(f"ROUGE-L: {rouge_l_score['Rouge-L-F']:.4f}")
179 | 
180 |     # meteor
181 |     meteor = corpus_meteor(combined_refs, combined_txts)
182 |     print(f"METEOR: {meteor:.4f}")
183 | 
184 |     # style CAV (w/ domain control)
185 |     acc = torch.zeros(1)
186 |     total = 0
187 | 
188 |     for author_id in tqdm(author_ids, disable=True):
189 |         if author_id in dc_cav_samples:
190 |             a_vecs = []
191 |             p_vecs = []
192 |             n_vecs = []
193 |             author_samples = dc_cav_samples[author_id]
194 |             for a_id, p_id, neg_author, n_id, _, _ in author_samples:
195 |                 assert neg_author != author_id
196 |                 a_vecs.append(gen_style_vecs[author_id][a_id])
197 |                 p_vecs.append(label_style_vecs[author_id][p_id])
198 |                 n_vecs.append(label_style_vecs[neg_author][n_id])
199 |             a_vecs = torch.vstack(a_vecs)
200 |             p_vecs = torch.vstack(p_vecs)
201 |             n_vecs = torch.vstack(n_vecs)
202 | 
203 |             sp_cossim = cossim(a_vecs, p_vecs)
204 |             sn_cossim = cossim(a_vecs, n_vecs)
205 |             acc += (sp_cossim > sn_cossim).float().sum()
206 |             
207 |             total += len(author_samples)
208 |     style_acc = float(acc / total)
209 | 
210 |     # semantic CAV (w/ domain control)
211 |     acc = torch.zeros(1)
212 |     total = 0
213 | 
214 |     for author_id in tqdm(author_ids, disable=True):
215 |         if author_id in dc_cav_samples:
216 |             a_vecs = []
217 |             p_vecs = []
218 |             n_vecs = []
219 |             author_samples = dc_cav_samples[author_id]
220 |             for a_id, p_id, neg_author, n_id, _, _ in author_samples:
221 |                 assert neg_author != author_id
222 |                 a_vecs.append(gen_semantic_vecs[author_id][a_id])
223 |                 p_vecs.append(label_semantic_vecs[author_id][p_id])
224 |                 n_vecs.append(label_semantic_vecs[neg_author][n_id])
225 |             a_vecs = torch.vstack(a_vecs)
226 |             p_vecs = torch.vstack(p_vecs)
227 |             n_vecs = torch.vstack(n_vecs)
228 | 
229 |             sp_cossim = cossim(a_vecs, p_vecs)
230 |             sn_cossim = cossim(a_vecs, n_vecs)
231 |             acc += (sp_cossim > sn_cossim).float().sum()
232 |             
233 |             total += len(author_samples)
234 |     semantic_acc = float(acc / total)
235 | 
236 |     print(f"CAV accuracy (style | semantic): {style_acc:.4f} | {semantic_acc:.4f}")
237 | 
238 |     # style embedding similarity
239 |     sum_cossim = torch.zeros(1)
240 | 
241 |     for author_id in tqdm(author_ids, disable=True):
242 |         mean_cossim = cossim(torch.vstack(label_style_vecs[author_id]), torch.vstack(gen_style_vecs[author_id])).mean()
243 |         sum_cossim += mean_cossim
244 |     style_cossim = float(sum_cossim / len(author_ids))
245 |     
246 |     # semantic embedding similarity
247 |     sum_cossim = torch.zeros(1)
248 | 
249 |     for author_id in tqdm(author_ids, disable=True):
250 |         mean_cossim = cossim(torch.vstack(label_semantic_vecs[author_id]), torch.vstack(gen_semantic_vecs[author_id])).mean()
251 |         sum_cossim += mean_cossim
252 |     semantic_cossim = float(sum_cossim / len(author_ids))
253 | 
254 |     print(f"Embedding Similarity (style | semantic): {style_cossim:.4f} | {semantic_cossim:.4f}")


--------------------------------------------------------------------------------
/src/utils/generator_utils.py:
--------------------------------------------------------------------------------
  1 | from tqdm.auto import tqdm
  2 | from dataclasses import dataclass
  3 | from typing import Any, Callable, Dict, List, NewType, Optional, Tuple, Union
  4 | 
  5 | import torch
  6 | from transformers import PreTrainedTokenizerBase
  7 | from transformers.file_utils import PaddingStrategy
  8 | 
  9 | 
 10 | def encode_example(example, tokenizer, roberta_tokenizer):
 11 |     max_len=128
 12 |     max_ref_len = {
 13 |         "style": 64,
 14 |         "semantic": 64,
 15 |     }
 16 | 
 17 |     src = tokenizer.eos_token + example['src'].replace("<|TITLE|> ", "").replace(" <|EOS|> ", tokenizer.eos_token)
 18 |     tgt = tokenizer.eos_token + example['tgt'] + tokenizer.eos_token
 19 |     srctgt = src + tgt + tokenizer.eos_token
 20 |     author = example.get('author_id', -1)
 21 |     assert isinstance(author, int)
 22 |     encoded_src = tokenizer(src)
 23 |     encoded_tgt = tokenizer(tgt)
 24 |     src_ids = encoded_src.input_ids
 25 |     tgt_ids = encoded_tgt.input_ids
 26 | 
 27 |     encoded_context = roberta_tokenizer(
 28 |         example['src'].replace("<|TITLE|> ", "").replace(" <|EOS|> ", roberta_tokenizer.eos_token)
 29 |     )
 30 |     context_ids = encoded_context.input_ids[-max_len // 2:]
 31 |     context_attention_mask = [1] * len(context_ids)
 32 | 
 33 |     srctgt_ids = src_ids + tgt_ids
 34 |     if (len(src_ids) + len(tgt_ids)) > max_len:
 35 |         if len(tgt_ids) > max_len:
 36 |             srctgt_ids = tgt_ids[:max_len]
 37 |         else:
 38 |             srctgt_ids = srctgt_ids[-max_len:]
 39 |     label_len = min(max_len, len(tgt_ids))
 40 |     input_ids = srctgt_ids
 41 |     attention_mask = [1] * len(input_ids)
 42 |     labels = [-100] * (len(srctgt_ids) - label_len) + tgt_ids[:max_len]
 43 |     assert len(attention_mask) == len(input_ids)
 44 |     assert len(labels) == len(input_ids)
 45 | 
 46 |     encoded_example = {
 47 |         "input_ids": input_ids,
 48 |         "attention_mask": attention_mask,
 49 |         "labels": labels,
 50 |         "author_labels": author,
 51 |         "context_ids": context_ids,
 52 |         "context_attention_mask": context_attention_mask,
 53 |     }
 54 | 
 55 |     ref_src = {
 56 |         "style": example['style_tgts'],
 57 |         "semantic": example['semantic_tgts'],
 58 |     }
 59 |     for ref_type, ref in ref_src.items():
 60 |         if ref is not None:
 61 |             # tokenize with roberta tokenizer
 62 |             encoded_ref = [roberta_tokenizer(item) for item in ref]
 63 |             ref_ids = [item.input_ids for item in encoded_ref]
 64 |             ref_attention_mask = [item.attention_mask for item in encoded_ref]
 65 |             for i in range(len(ref_ids)):
 66 |                 if max_ref_len[ref_type] and len(ref_ids[i]) > max_ref_len[ref_type]:
 67 |                     ref_ids[i] = [ref_ids[i][0]] + ref_ids[i][1:max_ref_len[ref_type]-1] + [ref_ids[i][-1]]
 68 |                     ref_attention_mask[i] = [ref_attention_mask[i][0]] + ref_attention_mask[i][1:max_ref_len[ref_type]-1] + [ref_attention_mask[i][-1]]
 69 |             encoded_example[f"{ref_type}_ids"] = ref_ids
 70 |             encoded_example[f"{ref_type}_attention_mask"] = ref_attention_mask
 71 | 
 72 |             # tokenize with dialogpt tokenizer
 73 |             encoded_ref = [tokenizer(tokenizer.eos_token + item) for item in ref]
 74 |             ref_ids = [item.input_ids for item in encoded_ref]
 75 |             ref_attention_mask = [item.attention_mask for item in encoded_ref]
 76 |             for i in range(len(ref_ids)):
 77 |                 if max_ref_len[ref_type] and len(ref_ids[i]) > max_ref_len[ref_type]:
 78 |                     ref_ids[i] = ref_ids[i][:max_ref_len[ref_type]]
 79 |                     ref_attention_mask[i] = ref_attention_mask[i][:max_ref_len[ref_type]]
 80 |             encoded_example[f"decoder_{ref_type}_ids"] = ref_ids
 81 |             encoded_example[f"decoder_{ref_type}_attention_mask"] = ref_attention_mask
 82 | 
 83 |     return encoded_example
 84 | 
 85 | 
 86 | @dataclass
 87 | class DataCollator:
 88 |     """
 89 |     Data collator that will dynamically pad the inputs received, as well as the labels.
 90 |     Args:
 91 |         tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):
 92 |             The tokenizer used for encoding the data.
 93 |         model ([`PreTrainedModel`]):
 94 |             The model that is being trained. If set and has the *prepare_decoder_input_ids_from_labels*, use it to
 95 |             prepare the *decoder_input_ids*
 96 |             This is useful when using *label_smoothing* to avoid calculating loss twice.
 97 |         padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):
 98 |             Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
 99 |             among:
100 |             - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single sequence
101 |               is provided).
102 |             - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
103 |               acceptable input length for the model if that argument is not provided.
104 |             - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
105 |               lengths).
106 |         max_length (`int`, *optional*):
107 |             Maximum length of the returned list and optionally padding length (see above).
108 |         pad_to_multiple_of (`int`, *optional*):
109 |             If set will pad the sequence to a multiple of the provided value.
110 |             This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
111 |             7.5 (Volta).
112 |         label_pad_token_id (`int`, *optional*, defaults to -100):
113 |             The id to use when padding the labels (-100 will be automatically ignored by PyTorch loss functions).
114 |         return_tensors (`str`):
115 |             The type of Tensor to return. Allowable values are "np", "pt" and "tf".
116 |     """
117 | 
118 |     tokenizer: PreTrainedTokenizerBase
119 |     roberta_tokenizer: PreTrainedTokenizerBase
120 |     model: Optional[Any] = None
121 |     padding: Union[bool, str, PaddingStrategy] = True
122 |     max_length: Optional[int] = None
123 |     pad_to_multiple_of: Optional[int] = None
124 |     label_pad_token_id: int = -100
125 |     return_tensors: str = "pt"
126 |     max_style_num: int = 10
127 |     max_semantic_num: int = 10
128 | 
129 |     def __call__(self, features, return_tensors=None):
130 |         import numpy as np
131 | 
132 |         if return_tensors is None:
133 |             return_tensors = self.return_tensors
134 | 
135 |         max_ref_num = {
136 |             "style": self.max_style_num,
137 |             "semantic": self.max_semantic_num,
138 |         }
139 | 
140 |         decoder_ref_type = None
141 |         if "decoder_style_ids" in features[0].keys():
142 |             assert not "decoder_semantic_ids" in features[0].keys()
143 |             decoder_ref_type = 'style'
144 |         elif "decoder_semantic_ids" in features[0].keys():
145 |             decoder_ref_type = 'semantic'
146 |         if decoder_ref_type is not None:
147 |             for feature in features:
148 |                 feature[f"decoder_{decoder_ref_type}_ids"] = np.concatenate(feature[f"decoder_{decoder_ref_type}_ids"][:max_ref_num[decoder_ref_type]])
149 |                 feature[f"decoder_{decoder_ref_type}_attention_mask"] = np.concatenate(feature[f"decoder_{decoder_ref_type}_attention_mask"][:max_ref_num[decoder_ref_type]])
150 |                 feature["input_ids"] = np.concatenate([feature[f"decoder_{decoder_ref_type}_ids"], feature["input_ids"]])
151 |                 feature["attention_mask"] = np.concatenate([feature[f"decoder_{decoder_ref_type}_attention_mask"], feature["attention_mask"]])
152 |                 if "labels" in features[0].keys():
153 |                     feature["labels"] = np.concatenate([[-100] * len(feature[f"decoder_{decoder_ref_type}_ids"]), feature["labels"]])
154 |                 del feature[f"decoder_{decoder_ref_type}_ids"]
155 |                 del feature[f"decoder_{decoder_ref_type}_attention_mask"]
156 |     
157 |         labels = [feature["labels"] for feature in features] if "labels" in features[0].keys() else None
158 |         context_ids = [feature["context_ids"] for feature in features] if "context_ids" in features[0].keys() else None
159 |         ref_ids = {
160 |             ref_type: 
161 |             [feature[f"{ref_type}_ids"] for feature in features] if f"{ref_type}_ids" in features[0].keys() else None
162 |             for ref_type in ["style", "semantic"]
163 |         }
164 |         # We have to pad the labels before calling `tokenizer.pad` as this method won't pad them and needs them of the
165 |         # same length to return tensors.
166 |         if labels is not None:
167 |             max_label_length = max(len(l) for l in labels)
168 |             if self.pad_to_multiple_of is not None:
169 |                 max_label_length = (
170 |                     (max_label_length + self.pad_to_multiple_of - 1)
171 |                     // self.pad_to_multiple_of
172 |                     * self.pad_to_multiple_of
173 |                 )
174 | 
175 |             padding_side = self.tokenizer.padding_side
176 |             assert padding_side == 'right'
177 |             for feature in features:
178 |                 remainder = [self.label_pad_token_id] * (max_label_length - len(feature["labels"]))
179 |                 if isinstance(feature["labels"], list):
180 |                     feature["labels"] = (
181 |                         feature["labels"] + remainder if padding_side == "right" else remainder + feature["labels"]
182 |                     )
183 |                 elif padding_side == "right":
184 |                     feature["labels"] = np.concatenate([feature["labels"], remainder]).astype(np.int64)
185 |                 else:
186 |                     feature["labels"] = np.concatenate([remainder, feature["labels"]]).astype(np.int64)
187 | 
188 |         if context_ids is not None:
189 |             max_context_length = max(len(c) for c in context_ids)
190 |             for feature in features:
191 |                 remainder = [self.roberta_tokenizer.pad_token_id] * (max_context_length - len(feature["context_ids"]))
192 |                 attn_mask_remainder = [0] * (max_context_length - len(feature["context_attention_mask"]))
193 |                 feature["context_ids"] = np.concatenate([feature["context_ids"], remainder]).astype(np.int64)
194 |                 feature["context_attention_mask"] = np.concatenate([feature["context_attention_mask"], attn_mask_remainder]).astype(np.int64)
195 | 
196 |         # get max length of all ref
197 |         max_ref_length = 0
198 |         for ref_type, type_ids in ref_ids.items():
199 |             if type_ids is None:
200 |                 continue
201 |             max_type_num = min(max(len(hs) for hs in type_ids), max_ref_num[ref_type])
202 |             max_type_length = max(len(h) for hs in type_ids for type_i, h in enumerate(hs) if type_i <= max_type_num)
203 |             max_ref_length = max(max_ref_length, max_type_length)
204 | 
205 |         for ref_type, type_ids in ref_ids.items():
206 |             if type_ids is None:
207 |                 continue
208 |             max_type_num = min(max(len(hs) for hs in type_ids), max_ref_num[ref_type])
209 |             max_type_length = max_ref_length
210 | 
211 |             for feature in features:
212 |                 if len(feature[f"{ref_type}_ids"]) < max_type_num:
213 |                     remainder = [[self.roberta_tokenizer.pad_token_id]] * (max_type_num - len(feature[f"{ref_type}_ids"]))
214 |                     attn_mask_remainder = [[0]] * (max_type_num - len(feature[f"{ref_type}_ids"]))
215 |                     if isinstance(feature[f"{ref_type}_ids"], list):
216 |                         feature[f"{ref_type}_ids"] += remainder
217 |                         feature[f"{ref_type}_attention_mask"] += attn_mask_remainder
218 |                     else:
219 |                         assert False
220 |                 if len(feature[f"{ref_type}_ids"]) > max_type_num:
221 |                     feature[f"{ref_type}_ids"] = feature[f"{ref_type}_ids"][:max_type_num]
222 |                     feature[f"{ref_type}_attention_mask"] = feature[f"{ref_type}_attention_mask"][:max_type_num]
223 |                 
224 |                 for type_i in range(max_type_num):
225 |                     remainder = [self.roberta_tokenizer.pad_token_id] * (max_type_length - len(feature[f"{ref_type}_ids"][type_i]))
226 |                     attn_mask_remainder = [0] * (max_type_length - len(feature[f"{ref_type}_attention_mask"][type_i]))
227 |                     if isinstance(feature[f"{ref_type}_ids"][type_i], list):
228 |                         feature[f"{ref_type}_ids"][type_i] += remainder
229 |                         feature[f"{ref_type}_attention_mask"][type_i] += attn_mask_remainder
230 |                     else:
231 |                         feature[f"{ref_type}_ids"][type_i] = np.concatenate(
232 |                             [feature[f"{ref_type}_ids"][type_i], remainder]
233 |                         ).astype(np.int64)
234 |                         feature[f"{ref_type}_attention_mask"][type_i] = np.concatenate(
235 |                             [feature[f"{ref_type}_attention_mask"][type_i], attn_mask_remainder]
236 |                         ).astype(np.int64)
237 | 
238 |         features = self.tokenizer.pad(
239 |             features,
240 |             padding=self.padding,
241 |             max_length=self.max_length,
242 |             pad_to_multiple_of=self.pad_to_multiple_of,
243 |             return_tensors=return_tensors,
244 |         )
245 | 
246 |         if "style_ids" in features and "semantic_ids" in features:
247 |             features["ref_input_ids"] = torch.cat((features["style_ids"][:, :5, :], features["semantic_ids"][:, :5, :]), 1)
248 |             features["ref_attention_mask"] = torch.cat((features["style_attention_mask"][:, :5, :], features["semantic_attention_mask"][:, :5, :]), 1)
249 |         else:
250 |             ref_type = None
251 |             if "style_ids" in features:
252 |                 ref_type = 'style'
253 |             elif "semantic_ids" in features:
254 |                 ref_type = 'semantic'
255 |             if ref_type is not None:
256 |                 features["ref_input_ids"] = features[f"{ref_type}_ids"]
257 |                 features["ref_attention_mask"] = features[f"{ref_type}_attention_mask"]
258 |         
259 |         if f"context_ids" in features:
260 |             features["context_input_ids"] = features["context_ids"]
261 |             del features["context_ids"]
262 | 
263 |         for key in list(features.keys()):
264 |             if (key.startswith("style") or key.startswith("semantic")) and "ref" not in key:
265 |                 del features[key]
266 | 
267 |         return features
268 | 
269 | 
270 | def generate_responses(trainer,
271 |                        dataset,
272 |                        max_length=128,
273 |                        disable_tqdm=False,
274 |                        out_file=None,
275 |                        **kwargs):
276 |     
277 |     trainer.model.eval()
278 | 
279 |     original_per_device_eval_batch_size = trainer.args.per_device_eval_batch_size
280 |     trainer.args.per_device_eval_batch_size = original_per_device_eval_batch_size // kwargs.get("num_beams", 1)
281 | 
282 |     tokenizer = trainer.tokenizer
283 | 
284 |     ids = []
285 |     authors = []
286 |     input_strs = []
287 |     label_strs = []
288 |     generated_strs = []
289 |     
290 |     for batch in tqdm(trainer.get_test_dataloader(dataset), disable=disable_tqdm):
291 |         batch_ids = batch['id'].tolist()
292 |         batch_authors = batch['author_id'].tolist()
293 |         del batch['id']
294 |         del batch['author_id']
295 |         batch_ids = [-1] * len(batch['input_ids'])
296 |         batch_authors = [-1] * len(batch['input_ids'])
297 |         batch_labels = batch['labels']
298 |         
299 |         batch_gen_input_ids = []
300 |         batch_gen_attention_mask = []
301 |         max_input_len = 0
302 |         
303 |         for input_ids, attention_mask, labels in zip(batch['input_ids'], batch['attention_mask'], batch['labels']):
304 |             gen_input_ids = input_ids[(attention_mask == 1) & (labels == -100)].tolist() + [50256]
305 |             gen_attention_mask = [1] * len(gen_input_ids)
306 |             batch_gen_input_ids.append(gen_input_ids)
307 |             batch_gen_attention_mask.append(gen_attention_mask)
308 |             max_input_len = max(max_input_len, len(gen_input_ids))
309 |         for i in range(len(batch_gen_input_ids)):
310 |             input_len = len(batch_gen_input_ids[i])
311 |             if input_len < max_input_len:
312 |                 batch_gen_input_ids[i] = [50256] * (max_input_len - input_len) + batch_gen_input_ids[i]
313 |                 batch_gen_attention_mask[i] = [0] * (max_input_len - input_len) + batch_gen_attention_mask[i]
314 |         batch['input_ids'] = torch.tensor(batch_gen_input_ids)
315 |         batch['attention_mask'] = torch.tensor(batch_gen_attention_mask)
316 |         
317 |         batch_input_strs = tokenizer.batch_decode(batch['input_ids'], skip_special_tokens=False)
318 |         batch_input_strs = [item.replace('<|endoftext|>', '\n').strip() for item in batch_input_strs]
319 |             
320 |         batch_labels[batch_labels == -100] = tokenizer.pad_token_id
321 |         
322 |         batch_label_strs = tokenizer.batch_decode(batch_labels, skip_special_tokens=True)
323 |         batch = {arg: val.to(trainer.args.device) if not isinstance(val, list) else [item.to(trainer.args.device) for item in val]
324 |                 for arg, val in batch.items()
325 |                 if arg not in ['author_id', 'labels', 'decoder_input_ids']}
326 |         assert max_input_len == batch['input_ids'].size(1)
327 |         batch_generated = trainer.model.generate(**batch, 
328 |                                                  max_length=max_length + max_input_len,
329 |                                                  pad_token_id=tokenizer.eos_token_id,
330 |                                                  **kwargs)
331 |         batch_generated = batch_generated[:, max_input_len:]
332 |         batch_generated_strs = tokenizer.batch_decode(batch_generated, skip_special_tokens=True)
333 |         batch_generated_strs = [batch_generated_strs[i*1:(i+1)*1] 
334 |                                 for i in range(len(batch_generated_strs)//1)]
335 | 
336 |         authors += batch_authors
337 |         input_strs += batch_input_strs
338 |         label_strs += batch_label_strs
339 |         generated_strs += batch_generated_strs
340 | 
341 |     generated_results = []
342 |     for idx, author, input_str, label_str, generated_str in zip(ids, authors, input_strs, label_strs, generated_strs):
343 |         generated_results.append(
344 |             {
345 |                 "id" : idx,
346 |                 "author": author,
347 |                 "input": input_str,
348 |                 "label": label_str,
349 |                 "generated": generated_str
350 |             }
351 |         )
352 |         
353 |     trainer.args.per_device_eval_batch_size = original_per_device_eval_batch_size
354 | 
355 |     return generated_results


--------------------------------------------------------------------------------
/src/models/dialogpt.py:
--------------------------------------------------------------------------------
   1 | # coding=utf-8
   2 | # Copyright 2018 The OpenAI Team Authors and HuggingFace Inc. team.
   3 | # Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
   4 | #
   5 | # Licensed under the Apache License, Version 2.0 (the "License");
   6 | # you may not use this file except in compliance with the License.
   7 | # You may obtain a copy of the License at
   8 | #
   9 | #     http://www.apache.org/licenses/LICENSE-2.0
  10 | #
  11 | # Unless required by applicable law or agreed to in writing, software
  12 | # distributed under the License is distributed on an "AS IS" BASIS,
  13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14 | # See the License for the specific language governing permissions and
  15 | # limitations under the License.
  16 | """PyTorch OpenAI GPT-2 model."""
  17 | 
  18 | import math
  19 | import os
  20 | from dataclasses import dataclass
  21 | from typing import Optional, Tuple, Dict, Union, Any
  22 | 
  23 | import torch
  24 | import torch.utils.checkpoint
  25 | from packaging import version
  26 | from torch import nn
  27 | from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  28 | 
  29 | 
  30 | if version.parse(torch.__version__) >= version.parse("1.6"):
  31 |     is_amp_available = True
  32 |     from torch.cuda.amp import autocast
  33 | else:
  34 |     is_amp_available = False
  35 | 
  36 | from transformers.activations import ACT2FN
  37 | from transformers.file_utils import (
  38 |     ModelOutput,
  39 |     add_code_sample_docstrings,
  40 |     add_start_docstrings,
  41 |     add_start_docstrings_to_model_forward,
  42 |     replace_return_docstrings,
  43 | )
  44 | from transformers.modeling_outputs import (
  45 |     BaseModelOutputWithPastAndCrossAttentions,
  46 |     CausalLMOutputWithCrossAttentions,
  47 |     SequenceClassifierOutputWithPast,
  48 |     TokenClassifierOutput,
  49 | )
  50 | from transformers.modeling_utils import (
  51 |     Conv1D,
  52 |     PreTrainedModel,
  53 |     SequenceSummary,
  54 |     find_pruneable_heads_and_indices,
  55 |     prune_conv1d_layer,
  56 | )
  57 | from transformers import logging
  58 | from transformers.utils.model_parallel_utils import assert_device_map, get_device_map
  59 | from transformers import GPT2Config
  60 | from transformers import AutoModel, AutoTokenizer
  61 | 
  62 | ### start of import statements for generation ###
  63 | import inspect
  64 | import warnings
  65 | from dataclasses import dataclass
  66 | from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
  67 | 
  68 | import torch
  69 | import torch.distributed as dist
  70 | from torch import nn
  71 | 
  72 | from transformers.file_utils import ModelOutput
  73 | from transformers.generation_beam_constraints import Constraint
  74 | from transformers.generation_beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
  75 | from transformers.generation_logits_process import (
  76 |     EncoderNoRepeatNGramLogitsProcessor,
  77 |     ForcedBOSTokenLogitsProcessor,
  78 |     ForcedEOSTokenLogitsProcessor,
  79 |     HammingDiversityLogitsProcessor,
  80 |     InfNanRemoveLogitsProcessor,
  81 |     LogitsProcessorList,
  82 |     MinLengthLogitsProcessor,
  83 |     NoBadWordsLogitsProcessor,
  84 |     NoRepeatNGramLogitsProcessor,
  85 |     PrefixConstrainedLogitsProcessor,
  86 |     RepetitionPenaltyLogitsProcessor,
  87 |     TemperatureLogitsWarper,
  88 |     TopKLogitsWarper,
  89 |     TopPLogitsWarper,
  90 |     TypicalLogitsWarper,
  91 | )
  92 | from transformers.generation_stopping_criteria import (
  93 |     MaxLengthCriteria,
  94 |     MaxTimeCriteria,
  95 |     StoppingCriteria,
  96 |     StoppingCriteriaList,
  97 |     validate_stopping_criteria,
  98 | )
  99 | from transformers.pytorch_utils import torch_int_div
 100 | from transformers.utils import logging
 101 | 
 102 | from transformers.generation_utils import (
 103 |     GreedySearchOutput,
 104 |     SampleOutput,
 105 |     BeamSearchOutput,
 106 |     BeamSampleOutput,
 107 |     GreedySearchEncoderDecoderOutput,
 108 | )
 109 | ### end of import statements for generation ###
 110 | 
 111 | logger = logging.get_logger(__name__)
 112 | 
 113 | _CHECKPOINT_FOR_DOC = "gpt2"
 114 | _CONFIG_FOR_DOC = "GPT2Config"
 115 | _TOKENIZER_FOR_DOC = "GPT2Tokenizer"
 116 | 
 117 | GPT2_PRETRAINED_MODEL_ARCHIVE_LIST = [
 118 |     "gpt2",
 119 |     "gpt2-medium",
 120 |     "gpt2-large",
 121 |     "gpt2-xl",
 122 |     "distilgpt2",
 123 |     # See all GPT-2 models at https://huggingface.co/models?filter=gpt2
 124 | ]
 125 | 
 126 | 
 127 | def load_tf_weights_in_gpt2(model, config, gpt2_checkpoint_path):
 128 |     """Load tf checkpoints in a pytorch model"""
 129 |     try:
 130 |         import re
 131 | 
 132 |         import tensorflow as tf
 133 |     except ImportError:
 134 |         logger.error(
 135 |             "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
 136 |             "https://www.tensorflow.org/install/ for installation instructions."
 137 |         )
 138 |         raise
 139 |     tf_path = os.path.abspath(gpt2_checkpoint_path)
 140 |     logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
 141 |     # Load weights from TF model
 142 |     init_vars = tf.train.list_variables(tf_path)
 143 |     names = []
 144 |     arrays = []
 145 |     for name, shape in init_vars:
 146 |         logger.info(f"Loading TF weight {name} with shape {shape}")
 147 |         array = tf.train.load_variable(tf_path, name)
 148 |         names.append(name)
 149 |         arrays.append(array.squeeze())
 150 | 
 151 |     for name, array in zip(names, arrays):
 152 |         name = name[6:]  # skip "model/"
 153 |         name = name.split("/")
 154 |         pointer = model
 155 |         for m_name in name:
 156 |             if re.fullmatch(r"[A-Za-z]+\d+", m_name):
 157 |                 scope_names = re.split(r"(\d+)", m_name)
 158 |             else:
 159 |                 scope_names = [m_name]
 160 |             if scope_names[0] == "w" or scope_names[0] == "g":
 161 |                 pointer = getattr(pointer, "weight")
 162 |             elif scope_names[0] == "b":
 163 |                 pointer = getattr(pointer, "bias")
 164 |             elif scope_names[0] == "wpe" or scope_names[0] == "wte":
 165 |                 pointer = getattr(pointer, scope_names[0])
 166 |                 pointer = getattr(pointer, "weight")
 167 |             else:
 168 |                 pointer = getattr(pointer, scope_names[0])
 169 |             if len(scope_names) >= 2:
 170 |                 num = int(scope_names[1])
 171 |                 pointer = pointer[num]
 172 |         try:
 173 |             assert (
 174 |                 pointer.shape == array.shape
 175 |             ), f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched"
 176 |         except AssertionError as e:
 177 |             e.args += (pointer.shape, array.shape)
 178 |             raise
 179 |         logger.info(f"Initialize PyTorch weight {name}")
 180 |         pointer.data = torch.from_numpy(array)
 181 |     return model
 182 | 
 183 | 
 184 | class GPT2Attention(nn.Module):
 185 |     def __init__(self, config, is_cross_attention=False, layer_idx=None):
 186 |         super().__init__()
 187 | 
 188 |         max_positions = config.max_position_embeddings
 189 |         self.register_buffer(
 190 |             "bias",
 191 |             torch.tril(torch.ones((max_positions, max_positions), dtype=torch.uint8)).view(
 192 |                 1, 1, max_positions, max_positions
 193 |             ),
 194 |         )
 195 |         self.register_buffer("masked_bias", torch.tensor(-1e4))
 196 | 
 197 |         self.embed_dim = config.hidden_size
 198 |         self.num_heads = config.num_attention_heads
 199 |         self.head_dim = self.embed_dim // self.num_heads
 200 |         self.split_size = self.embed_dim
 201 |         if self.head_dim * self.num_heads != self.embed_dim:
 202 |             raise ValueError(
 203 |                 f"`embed_dim` must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."
 204 |             )
 205 | 
 206 |         self.scale_attn_weights = config.scale_attn_weights
 207 |         self.is_cross_attention = is_cross_attention
 208 | 
 209 |         # Layer-wise attention scaling, reordering, and upcasting
 210 |         self.scale_attn_by_inverse_layer_idx = config.scale_attn_by_inverse_layer_idx
 211 |         self.layer_idx = layer_idx
 212 |         self.reorder_and_upcast_attn = config.reorder_and_upcast_attn
 213 | 
 214 |         if self.is_cross_attention:
 215 |             self.c_attn = Conv1D(2 * self.embed_dim, self.embed_dim)
 216 |             self.q_attn = Conv1D(self.embed_dim, self.embed_dim)
 217 |         else:
 218 |             self.c_attn = Conv1D(3 * self.embed_dim, self.embed_dim)
 219 |         self.c_proj = Conv1D(self.embed_dim, self.embed_dim)
 220 | 
 221 |         self.attn_dropout = nn.Dropout(config.attn_pdrop)
 222 |         self.resid_dropout = nn.Dropout(config.resid_pdrop)
 223 | 
 224 |         self.pruned_heads = set()
 225 | 
 226 |     def prune_heads(self, heads):
 227 |         if len(heads) == 0:
 228 |             return
 229 |         heads, index = find_pruneable_heads_and_indices(heads, self.num_heads, self.head_dim, self.pruned_heads)
 230 |         index_attn = torch.cat([index, index + self.split_size, index + (2 * self.split_size)])
 231 | 
 232 |         # Prune conv1d layers
 233 |         self.c_attn = prune_conv1d_layer(self.c_attn, index_attn, dim=1)
 234 |         self.c_proj = prune_conv1d_layer(self.c_proj, index, dim=0)
 235 | 
 236 |         # Update hyper params
 237 |         self.split_size = (self.split_size // self.num_heads) * (self.num_heads - len(heads))
 238 |         self.num_heads = self.num_heads - len(heads)
 239 |         self.pruned_heads = self.pruned_heads.union(heads)
 240 | 
 241 |     def _attn(self, query, key, value, attention_mask=None, head_mask=None):
 242 |         attn_weights = torch.matmul(query, key.transpose(-1, -2))
 243 | 
 244 |         if self.scale_attn_weights:
 245 |             attn_weights = attn_weights / (value.size(-1) ** 0.5)
 246 | 
 247 |         # Layer-wise attention scaling
 248 |         if self.scale_attn_by_inverse_layer_idx:
 249 |             attn_weights = attn_weights / float(self.layer_idx + 1)
 250 | 
 251 |         if not self.is_cross_attention:
 252 |             # if only "normal" attention layer implements causal mask
 253 |             query_length, key_length = query.size(-2), key.size(-2)
 254 |             causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].bool()
 255 |             attn_weights = torch.where(causal_mask, attn_weights, self.masked_bias.to(attn_weights.dtype))
 256 | 
 257 |         if attention_mask is not None:
 258 |             # Apply the attention mask
 259 |             attn_weights = attn_weights + attention_mask
 260 | 
 261 |         attn_weights = nn.functional.softmax(attn_weights, dim=-1)
 262 | 
 263 |         # Downcast (if necessary) back to V's dtype (if in mixed-precision) -- No-Op otherwise
 264 |         attn_weights = attn_weights.type(value.dtype)
 265 |         attn_weights = self.attn_dropout(attn_weights)
 266 | 
 267 |         # Mask heads if we want to
 268 |         if head_mask is not None:
 269 |             attn_weights = attn_weights * head_mask
 270 | 
 271 |         attn_output = torch.matmul(attn_weights, value)
 272 | 
 273 |         return attn_output, attn_weights
 274 | 
 275 |     def _upcast_and_reordered_attn(self, query, key, value, attention_mask=None, head_mask=None):
 276 |         # Use `torch.baddbmm` (a bit more efficient w/ alpha param for scaling -- from Megatron-LM)
 277 |         bsz, num_heads, q_seq_len, dk = query.size()
 278 |         _, _, k_seq_len, _ = key.size()
 279 | 
 280 |         # Preallocate attn_weights for `baddbmm`
 281 |         attn_weights = torch.empty(bsz * num_heads, q_seq_len, k_seq_len, dtype=torch.float32, device=query.device)
 282 | 
 283 |         # Compute Scale Factor
 284 |         scale_factor = 1.0
 285 |         if self.scale_attn_weights:
 286 |             scale_factor /= float(value.size(-1)) ** 0.5
 287 | 
 288 |         if self.scale_attn_by_inverse_layer_idx:
 289 |             scale_factor /= float(self.layer_idx + 1)
 290 | 
 291 |         # Upcast (turn off autocast) and reorder (Scale K by 1 / root(dk))
 292 |         if is_amp_available:
 293 |             with autocast(enabled=False):
 294 |                 q, k = query.reshape(-1, q_seq_len, dk), key.transpose(-1, -2).reshape(-1, dk, k_seq_len)
 295 |                 attn_weights = torch.baddbmm(attn_weights, q.float(), k.float(), beta=0, alpha=scale_factor)
 296 |                 attn_weights = attn_weights.reshape(bsz, num_heads, q_seq_len, k_seq_len)
 297 |         else:
 298 |             q, k = query.reshape(-1, q_seq_len, dk), key.transpose(-1, -2).reshape(-1, dk, k_seq_len)
 299 |             attn_weights = torch.baddbmm(attn_weights, q.float(), k.float(), beta=0, alpha=scale_factor)
 300 |             attn_weights = attn_weights.reshape(bsz, num_heads, q_seq_len, k_seq_len)
 301 | 
 302 |         if not self.is_cross_attention:
 303 |             # if only "normal" attention layer implements causal mask
 304 |             query_length, key_length = query.size(-2), key.size(-2)
 305 |             causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].bool()
 306 |             attn_weights = torch.where(causal_mask, attn_weights, self.masked_bias.to(attn_weights.dtype))
 307 | 
 308 |         if attention_mask is not None:
 309 |             # Apply the attention mask
 310 |             attn_weights = attn_weights + attention_mask
 311 | 
 312 |         attn_weights = nn.functional.softmax(attn_weights, dim=-1)
 313 | 
 314 |         # Downcast (if necessary) back to V's dtype (if in mixed-precision) -- No-Op if otherwise
 315 |         if attn_weights.dtype != torch.float32:
 316 |             raise RuntimeError("Error with upcasting, attn_weights does not have dtype torch.float32")
 317 |         attn_weights = attn_weights.type(value.dtype)
 318 |         attn_weights = self.attn_dropout(attn_weights)
 319 | 
 320 |         # Mask heads if we want to
 321 |         if head_mask is not None:
 322 |             attn_weights = attn_weights * head_mask
 323 | 
 324 |         attn_output = torch.matmul(attn_weights, value)
 325 | 
 326 |         return attn_output, attn_weights
 327 | 
 328 |     def _split_heads(self, tensor, num_heads, attn_head_size):
 329 |         """
 330 |         Splits hidden_size dim into attn_head_size and num_heads
 331 |         """
 332 |         new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
 333 |         tensor = tensor.view(new_shape)
 334 |         return tensor.permute(0, 2, 1, 3)  # (batch, head, seq_length, head_features)
 335 | 
 336 |     def _merge_heads(self, tensor, num_heads, attn_head_size):
 337 |         """
 338 |         Merges attn_head_size dim and num_attn_heads dim into hidden_size
 339 |         """
 340 |         tensor = tensor.permute(0, 2, 1, 3).contiguous()
 341 |         new_shape = tensor.size()[:-2] + (num_heads * attn_head_size,)
 342 |         return tensor.view(new_shape)
 343 | 
 344 |     def forward(
 345 |         self,
 346 |         hidden_states,
 347 |         layer_past=None,
 348 |         attention_mask=None,
 349 |         head_mask=None,
 350 |         encoder_hidden_states=None,
 351 |         encoder_attention_mask=None,
 352 |         use_cache=False,
 353 |         output_attentions=False,
 354 |     ):
 355 |         if encoder_hidden_states is not None:
 356 |             if not hasattr(self, "q_attn"):
 357 |                 raise ValueError(
 358 |                     "If class is used as cross attention, the weights `q_attn` have to be defined. "
 359 |                     "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`."
 360 |                 )
 361 | 
 362 |             query = self.q_attn(hidden_states)
 363 |             key, value = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
 364 |             attention_mask = encoder_attention_mask
 365 |         else:
 366 |             query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)
 367 | 
 368 |         query = self._split_heads(query, self.num_heads, self.head_dim)
 369 |         key = self._split_heads(key, self.num_heads, self.head_dim)
 370 |         value = self._split_heads(value, self.num_heads, self.head_dim)
 371 | 
 372 |         if layer_past is not None:
 373 |             past_key, past_value = layer_past
 374 |             key = torch.cat((past_key, key), dim=-2)
 375 |             value = torch.cat((past_value, value), dim=-2)
 376 | 
 377 |         if use_cache is True:
 378 |             present = (key, value)
 379 |         else:
 380 |             present = None
 381 | 
 382 |         if self.reorder_and_upcast_attn:
 383 |             attn_output, attn_weights = self._upcast_and_reordered_attn(query, key, value, attention_mask, head_mask)
 384 |         else:
 385 |             attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
 386 | 
 387 |         attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
 388 |         attn_output = self.c_proj(attn_output)
 389 |         attn_output = self.resid_dropout(attn_output)
 390 | 
 391 |         outputs = (attn_output, present)
 392 |         if output_attentions:
 393 |             outputs += (attn_weights,)
 394 | 
 395 |         return outputs  # a, present, (attentions)
 396 | 
 397 | 
 398 | class GPT2MLP(nn.Module):
 399 |     def __init__(self, intermediate_size, config):
 400 |         super().__init__()
 401 |         embed_dim = config.hidden_size
 402 |         self.c_fc = Conv1D(intermediate_size, embed_dim)
 403 |         self.c_proj = Conv1D(embed_dim, intermediate_size)
 404 |         self.act = ACT2FN[config.activation_function]
 405 |         self.dropout = nn.Dropout(config.resid_pdrop)
 406 | 
 407 |     def forward(self, hidden_states):
 408 |         hidden_states = self.c_fc(hidden_states)
 409 |         hidden_states = self.act(hidden_states)
 410 |         hidden_states = self.c_proj(hidden_states)
 411 |         hidden_states = self.dropout(hidden_states)
 412 |         return hidden_states
 413 | 
 414 | 
 415 | class GPT2Block(nn.Module):
 416 |     def __init__(self, config, layer_idx=None):
 417 |         super().__init__()
 418 |         hidden_size = config.hidden_size
 419 |         inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size
 420 | 
 421 |         self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
 422 |         self.attn = GPT2Attention(config, layer_idx=layer_idx)
 423 |         self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
 424 | 
 425 |         if config.add_cross_attention:
 426 |             self.crossattention = GPT2Attention(config, is_cross_attention=True, layer_idx=layer_idx)
 427 |             self.ln_cross_attn = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
 428 | 
 429 |         self.mlp = GPT2MLP(inner_dim, config)
 430 | 
 431 |     def forward(
 432 |         self,
 433 |         hidden_states,
 434 |         layer_past=None,
 435 |         attention_mask=None,
 436 |         head_mask=None,
 437 |         encoder_hidden_states=None,
 438 |         encoder_attention_mask=None,
 439 |         use_cache=False,
 440 |         output_attentions=False,
 441 |     ):
 442 |         residual = hidden_states
 443 |         hidden_states = self.ln_1(hidden_states)
 444 |         attn_outputs = self.attn(
 445 |             hidden_states,
 446 |             layer_past=layer_past,
 447 |             attention_mask=attention_mask,
 448 |             head_mask=head_mask,
 449 |             use_cache=use_cache,
 450 |             output_attentions=output_attentions,
 451 |         )
 452 |         attn_output = attn_outputs[0]  # output_attn: a, present, (attentions)
 453 |         outputs = attn_outputs[1:]
 454 |         # residual connection
 455 |         hidden_states = attn_output + residual
 456 | 
 457 |         if encoder_hidden_states is not None:
 458 |             # add one self-attention block for cross-attention
 459 |             if not hasattr(self, "crossattention"):
 460 |                 raise ValueError(
 461 |                     f"If `encoder_hidden_states` are passed, {self} has to be instantiated with "
 462 |                     "cross-attention layers by setting `config.add_cross_attention=True`"
 463 |                 )
 464 |             residual = hidden_states
 465 |             hidden_states = self.ln_cross_attn(hidden_states)
 466 |             cross_attn_outputs = self.crossattention(
 467 |                 hidden_states,
 468 |                 attention_mask=attention_mask,
 469 |                 head_mask=head_mask,
 470 |                 encoder_hidden_states=encoder_hidden_states,
 471 |                 encoder_attention_mask=encoder_attention_mask,
 472 |                 output_attentions=output_attentions,
 473 |             )
 474 |             attn_output = cross_attn_outputs[0]
 475 |             # residual connection
 476 |             hidden_states = residual + attn_output
 477 |             outputs = outputs + cross_attn_outputs[2:]  # add cross attentions if we output attention weights
 478 | 
 479 |         residual = hidden_states
 480 |         hidden_states = self.ln_2(hidden_states)
 481 |         feed_forward_hidden_states = self.mlp(hidden_states)
 482 |         # residual connection
 483 |         hidden_states = residual + feed_forward_hidden_states
 484 | 
 485 |         if use_cache:
 486 |             outputs = (hidden_states,) + outputs
 487 |         else:
 488 |             outputs = (hidden_states,) + outputs[1:]
 489 | 
 490 |         return outputs  # hidden_states, present, (attentions, cross_attentions)
 491 | 
 492 | 
 493 | class GPT2PreTrainedModel(PreTrainedModel):
 494 |     """
 495 |     An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
 496 |     models.
 497 |     """
 498 | 
 499 |     config_class = GPT2Config
 500 |     load_tf_weights = load_tf_weights_in_gpt2
 501 |     base_model_prefix = "transformer"
 502 |     is_parallelizable = True
 503 |     supports_gradient_checkpointing = True
 504 | 
 505 |     def __init__(self, *inputs, **kwargs):
 506 |         super().__init__(*inputs, **kwargs)
 507 | 
 508 |     def _init_weights(self, module):
 509 |         """Initialize the weights."""
 510 |         if isinstance(module, (nn.Linear, Conv1D)):
 511 |             # Slightly different from the TF version which uses truncated_normal for initialization
 512 |             # cf https://github.com/pytorch/pytorch/pull/5617
 513 |             module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
 514 |             if module.bias is not None:
 515 |                 module.bias.data.zero_()
 516 |         elif isinstance(module, nn.Embedding):
 517 |             module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
 518 |             if module.padding_idx is not None:
 519 |                 module.weight.data[module.padding_idx].zero_()
 520 |         elif isinstance(module, nn.LayerNorm):
 521 |             module.bias.data.zero_()
 522 |             module.weight.data.fill_(1.0)
 523 | 
 524 |         # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
 525 |         #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
 526 |         #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
 527 |         #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
 528 |         #
 529 |         # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
 530 |         for name, p in module.named_parameters():
 531 |             if "c_proj" in name and "weight" in name:
 532 |                 # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
 533 |                 p.data.normal_(mean=0.0, std=(self.config.initializer_range / math.sqrt(2 * self.config.n_layer)))
 534 | 
 535 |     def _set_gradient_checkpointing(self, module, value=False):
 536 |         if isinstance(module, GPT2Model):
 537 |             module.gradient_checkpointing = value
 538 | 
 539 | 
 540 | @dataclass
 541 | class GPT2DoubleHeadsModelOutput(ModelOutput):
 542 |     """
 543 |     Base class for outputs of models predicting if two sentences are consecutive or not.
 544 | 
 545 |     Args:
 546 |         loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
 547 |             Language modeling loss.
 548 |         mc_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mc_labels` is provided):
 549 |             Multiple choice classification loss.
 550 |         logits (`torch.FloatTensor` of shape `(batch_size, num_choices, sequence_length, config.vocab_size)`):
 551 |             Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
 552 |         mc_logits (`torch.FloatTensor` of shape `(batch_size, num_choices)`):
 553 |             Prediction scores of the multiple choice classification head (scores for each choice before SoftMax).
 554 |         past_key_values (`Tuple[Tuple[torch.Tensor]]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
 555 |             Tuple of length `config.n_layers`, containing tuples of tensors of shape `(batch_size, num_heads,
 556 |             sequence_length, embed_size_per_head)`).
 557 | 
 558 |             Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see
 559 |             `past_key_values` input) to speed up sequential decoding.
 560 |         hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
 561 |             Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
 562 |             shape `(batch_size, sequence_length, hidden_size)`.
 563 | 
 564 |             Hidden-states of the model at the output of each layer plus the initial embedding outputs.
 565 |         attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
 566 |             Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
 567 |             sequence_length)`.
 568 | 
 569 |             GPT2Attentions weights after the attention softmax, used to compute the weighted average in the
 570 |             self-attention heads.
 571 |     """
 572 | 
 573 |     loss: Optional[torch.FloatTensor] = None
 574 |     mc_loss: Optional[torch.FloatTensor] = None
 575 |     logits: torch.FloatTensor = None
 576 |     mc_logits: torch.FloatTensor = None
 577 |     past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
 578 |     hidden_states: Optional[Tuple[torch.FloatTensor]] = None
 579 |     attentions: Optional[Tuple[torch.FloatTensor]] = None
 580 | 
 581 | 
 582 | GPT2_START_DOCSTRING = r"""
 583 | 
 584 |     This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
 585 |     library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
 586 |     etc.)
 587 | 
 588 |     This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
 589 |     Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
 590 |     and behavior.
 591 | 
 592 |     Parameters:
 593 |         config ([`GPT2Config`]): Model configuration class with all the parameters of the model.
 594 |             Initializing with a config file does not load the weights associated with the model, only the
 595 |             configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
 596 | """
 597 | 
 598 | GPT2_INPUTS_DOCSTRING = r"""
 599 |     Args:
 600 |         input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
 601 |             `input_ids_length` = `sequence_length` if `past_key_values` is `None` else
 602 |             `past_key_values[0][0].shape[-2]` (`sequence_length` of input past key value states). Indices of input
 603 |             sequence tokens in the vocabulary.
 604 | 
 605 |             If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
 606 |             `input_ids`.
 607 | 
 608 |             Indices can be obtained using [`GPT2Tokenizer`]. See [`PreTrainedTokenizer.encode`] and
 609 |             [`PreTrainedTokenizer.__call__`] for details.
 610 | 
 611 |             [What are input IDs?](../glossary#input-ids)
 612 |         past_key_values (`Tuple[Tuple[torch.Tensor]]` of length `config.n_layers`):
 613 |             Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see
 614 |             `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have
 615 |             their past given to this model should not be passed as `input_ids` as they have already been computed.
 616 |         attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
 617 |             Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
 618 | 
 619 |             - 1 for tokens that are **not masked**,
 620 |             - 0 for tokens that are **masked**.
 621 | 
 622 |             [What are attention masks?](../glossary#attention-mask)
 623 |         token_type_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`, *optional*):
 624 |             Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
 625 |             1]`:
 626 | 
 627 |             - 0 corresponds to a *sentence A* token,
 628 |             - 1 corresponds to a *sentence B* token.
 629 | 
 630 |             [What are token type IDs?](../glossary#token-type-ids)
 631 |         position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
 632 |             Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
 633 |             config.max_position_embeddings - 1]`.
 634 | 
 635 |             [What are position IDs?](../glossary#position-ids)
 636 |         head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
 637 |             Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
 638 | 
 639 |             - 1 indicates the head is **not masked**,
 640 |             - 0 indicates the head is **masked**.
 641 | 
 642 |         inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
 643 |             Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
 644 |             is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
 645 |             model's internal embedding lookup matrix.
 646 | 
 647 |             If `past_key_values` is used, optionally only the last `inputs_embeds` have to be input (see
 648 |             `past_key_values`).
 649 |         use_cache (`bool`, *optional*):
 650 |             If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
 651 |             `past_key_values`).
 652 |         output_attentions (`bool`, *optional*):
 653 |             Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
 654 |             tensors for more detail.
 655 |         output_hidden_states (`bool`, *optional*):
 656 |             Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
 657 |             more detail.
 658 |         return_dict (`bool`, *optional*):
 659 |             Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
 660 | """
 661 | PARALLELIZE_DOCSTRING = r"""
 662 |     This is an experimental feature and is a subject to change at a moment's notice.
 663 | 
 664 |     Uses a device map to distribute attention modules of the model across several devices. If no device map is given,
 665 |     it will evenly distribute blocks across all devices.
 666 | 
 667 |     Args:
 668 |         device_map (`Dict[int, list]`, optional, defaults to None):
 669 |             A dictionary that maps attention modules to devices. Note that the embedding module and LMHead are always
 670 |             automatically mapped to the first device (for esoteric reasons). That means that the first device should
 671 |             have fewer attention modules mapped to it than other devices. For reference, the gpt2 models have the
 672 |             following number of attention modules:
 673 | 
 674 |                 - gpt2: 12
 675 |                 - gpt2-medium: 24
 676 |                 - gpt2-large: 36
 677 |                 - gpt2-xl: 48
 678 | 
 679 |     Example:
 680 | 
 681 |     ```python
 682 |     # Here is an example of a device map on a machine with 4 GPUs using gpt2-xl, which has a total of 48 attention modules:
 683 |     model = GPT2LMHeadModel.from_pretrained("gpt2-xl")
 684 |     device_map = {
 685 |         0: [0, 1, 2, 3, 4, 5, 6, 7, 8],
 686 |         1: [9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21],
 687 |         2: [22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34],
 688 |         3: [35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47],
 689 |     }
 690 |     model.parallelize(device_map)
 691 |     ```
 692 | """
 693 | DEPARALLELIZE_DOCSTRING = r"""
 694 |     Moves the model to cpu from a model parallel state.
 695 | 
 696 |     Example:
 697 | 
 698 |     ```python
 699 |     # On a 4 GPU machine with gpt2-large:
 700 |     model = GPT2LMHeadModel.from_pretrained("gpt2-large")
 701 |     device_map = {
 702 |         0: [0, 1, 2, 3, 4, 5, 6, 7],
 703 |         1: [8, 9, 10, 11, 12, 13, 14, 15],
 704 |         2: [16, 17, 18, 19, 20, 21, 22, 23],
 705 |         3: [24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35],
 706 |     }
 707 |     model.parallelize(device_map)  # Splits the model across several devices
 708 |     model.deparallelize()  # Put the model back on cpu and cleans memory by calling torch.cuda.empty_cache()
 709 |     ```
 710 | """
 711 | 
 712 | 
 713 | @add_start_docstrings(
 714 |     "The bare GPT2 Model transformer outputting raw hidden-states without any specific head on top.",
 715 |     GPT2_START_DOCSTRING,
 716 | )
 717 | class GPT2Model(GPT2PreTrainedModel):
 718 |     _keys_to_ignore_on_load_missing = ["attn.masked_bias"]
 719 | 
 720 |     def __init__(self, config):
 721 |         super().__init__(config)
 722 | 
 723 |         self.embed_dim = config.hidden_size
 724 | 
 725 |         self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
 726 |         self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim)
 727 | 
 728 |         self.drop = nn.Dropout(config.embd_pdrop)
 729 |         self.h = nn.ModuleList([GPT2Block(config, layer_idx=i) for i in range(config.num_hidden_layers)])
 730 |         self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
 731 | 
 732 |         # Model parallel
 733 |         self.model_parallel = False
 734 |         self.device_map = None
 735 |         self.gradient_checkpointing = False
 736 | 
 737 |         # Initialize weights and apply final processing
 738 |         self.post_init()
 739 | 
 740 |     @add_start_docstrings(PARALLELIZE_DOCSTRING)
 741 |     def parallelize(self, device_map=None):
 742 |         # Check validity of device_map
 743 |         self.device_map = (
 744 |             get_device_map(len(self.h), range(torch.cuda.device_count())) if device_map is None else device_map
 745 |         )
 746 |         assert_device_map(self.device_map, len(self.h))
 747 |         self.model_parallel = True
 748 |         self.first_device = "cpu" if "cpu" in self.device_map.keys() else "cuda:" + str(min(self.device_map.keys()))
 749 |         self.last_device = "cuda:" + str(max(self.device_map.keys()))
 750 |         self.wte = self.wte.to(self.first_device)
 751 |         self.wpe = self.wpe.to(self.first_device)
 752 |         # Load onto devices
 753 |         for k, v in self.device_map.items():
 754 |             for block in v:
 755 |                 cuda_device = "cuda:" + str(k)
 756 |                 self.h[block] = self.h[block].to(cuda_device)
 757 |         # ln_f to last
 758 |         self.ln_f = self.ln_f.to(self.last_device)
 759 | 
 760 |     @add_start_docstrings(DEPARALLELIZE_DOCSTRING)
 761 |     def deparallelize(self):
 762 |         self.model_parallel = False
 763 |         self.device_map = None
 764 |         self.first_device = "cpu"
 765 |         self.last_device = "cpu"
 766 |         self.wte = self.wte.to("cpu")
 767 |         self.wpe = self.wpe.to("cpu")
 768 |         for index in range(len(self.h)):
 769 |             self.h[index] = self.h[index].to("cpu")
 770 |         self.ln_f = self.ln_f.to("cpu")
 771 |         torch.cuda.empty_cache()
 772 | 
 773 |     def get_input_embeddings(self):
 774 |         return self.wte
 775 | 
 776 |     def set_input_embeddings(self, new_embeddings):
 777 |         self.wte = new_embeddings
 778 | 
 779 |     def _prune_heads(self, heads_to_prune):
 780 |         """
 781 |         Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
 782 |         """
 783 |         for layer, heads in heads_to_prune.items():
 784 |             self.h[layer].attn.prune_heads(heads)
 785 | 
 786 |     @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
 787 |     @add_code_sample_docstrings(
 788 |         processor_class=_TOKENIZER_FOR_DOC,
 789 |         checkpoint=_CHECKPOINT_FOR_DOC,
 790 |         output_type=BaseModelOutputWithPastAndCrossAttentions,
 791 |         config_class=_CONFIG_FOR_DOC,
 792 |     )
 793 |     def forward(
 794 |         self,
 795 |         input_ids=None,
 796 |         past_key_values=None,
 797 |         attention_mask=None,
 798 |         token_type_ids=None,
 799 |         position_ids=None,
 800 |         head_mask=None,
 801 |         inputs_embeds=None,
 802 |         encoder_hidden_states=None,
 803 |         encoder_attention_mask=None,
 804 |         use_cache=None,
 805 |         output_attentions=None,
 806 |         output_hidden_states=None,
 807 |         return_dict=None,
 808 |     ):
 809 |         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
 810 |         output_hidden_states = (
 811 |             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
 812 |         )
 813 |         use_cache = use_cache if use_cache is not None else self.config.use_cache
 814 |         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 815 | 
 816 |         if input_ids is not None and inputs_embeds is not None:
 817 |             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
 818 |         elif input_ids is not None:
 819 |             input_shape = input_ids.size()
 820 |             input_ids = input_ids.view(-1, input_shape[-1])
 821 |             batch_size = input_ids.shape[0]
 822 |         elif inputs_embeds is not None:
 823 |             input_shape = inputs_embeds.size()[:-1]
 824 |             batch_size = inputs_embeds.shape[0]
 825 |         else:
 826 |             raise ValueError("You have to specify either input_ids or inputs_embeds")
 827 | 
 828 |         device = input_ids.device if input_ids is not None else inputs_embeds.device
 829 | 
 830 |         if token_type_ids is not None:
 831 |             token_type_ids = token_type_ids.view(-1, input_shape[-1])
 832 |         if position_ids is not None:
 833 |             position_ids = position_ids.view(-1, input_shape[-1])
 834 | 
 835 |         if past_key_values is None:
 836 |             past_length = 0
 837 |             past_key_values = tuple([None] * len(self.h))
 838 |         else:
 839 |             past_length = past_key_values[0][0].size(-2)
 840 |         if position_ids is None:
 841 |             position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
 842 |             position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])
 843 | 
 844 |         # GPT2Attention mask.
 845 |         if attention_mask is not None:
 846 |             if batch_size <= 0:
 847 |                 raise ValueError("batch_size has to be defined and > 0")
 848 |             attention_mask = attention_mask.view(batch_size, -1)
 849 |             # We create a 3D attention mask from a 2D tensor mask.
 850 |             # Sizes are [batch_size, 1, 1, to_seq_length]
 851 |             # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
 852 |             # this attention mask is more simple than the triangular masking of causal attention
 853 |             # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
 854 |             attention_mask = attention_mask[:, None, None, :]
 855 | 
 856 |             # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
 857 |             # masked positions, this operation will create a tensor which is 0.0 for
 858 |             # positions we want to attend and -10000.0 for masked positions.
 859 |             # Since we are adding it to the raw scores before the softmax, this is
 860 |             # effectively the same as removing these entirely.
 861 |             attention_mask = attention_mask.to(dtype=self.dtype)  # fp16 compatibility
 862 |             attention_mask = (1.0 - attention_mask) * -10000.0
 863 | 
 864 |         # If a 2D or 3D attention mask is provided for the cross-attention
 865 |         # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
 866 |         if self.config.add_cross_attention and encoder_hidden_states is not None:
 867 |             encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
 868 |             encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
 869 |             if encoder_attention_mask is None:
 870 |                 encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
 871 |             encoder_attention_mask = self.invert_attention_mask(encoder_attention_mask)
 872 |         else:
 873 |             encoder_attention_mask = None
 874 | 
 875 |         # Prepare head mask if needed
 876 |         # 1.0 in head_mask indicate we keep the head
 877 |         # attention_probs has shape bsz x n_heads x N x N
 878 |         # head_mask has shape n_layer x batch x n_heads x N x N
 879 |         head_mask = self.get_head_mask(head_mask, self.config.n_layer)
 880 | 
 881 |         if inputs_embeds is None:
 882 |             inputs_embeds = self.wte(input_ids)
 883 |         position_embeds = self.wpe(position_ids)
 884 |         hidden_states = inputs_embeds + position_embeds
 885 | 
 886 |         if token_type_ids is not None:
 887 |             token_type_embeds = self.wte(token_type_ids)
 888 |             hidden_states = hidden_states + token_type_embeds
 889 | 
 890 |         hidden_states = self.drop(hidden_states)
 891 | 
 892 |         output_shape = input_shape + (hidden_states.size(-1),)
 893 | 
 894 |         presents = () if use_cache else None
 895 |         all_self_attentions = () if output_attentions else None
 896 |         all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
 897 |         all_hidden_states = () if output_hidden_states else None
 898 |         for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
 899 | 
 900 |             # Model parallel
 901 |             if self.model_parallel:
 902 |                 torch.cuda.set_device(hidden_states.device)
 903 |                 # Ensure layer_past is on same device as hidden_states (might not be correct)
 904 |                 if layer_past is not None:
 905 |                     layer_past = tuple(past_state.to(hidden_states.device) for past_state in layer_past)
 906 |                 # Ensure that attention_mask is always on the same device as hidden_states
 907 |                 if attention_mask is not None:
 908 |                     attention_mask = attention_mask.to(hidden_states.device)
 909 |                 if isinstance(head_mask, torch.Tensor):
 910 |                     head_mask = head_mask.to(hidden_states.device)
 911 |             if output_hidden_states:
 912 |                 all_hidden_states = all_hidden_states + (hidden_states,)
 913 | 
 914 |             if self.gradient_checkpointing and self.training:
 915 | 
 916 |                 if use_cache:
 917 |                     logger.warning(
 918 |                         "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
 919 |                     )
 920 |                     use_cache = False
 921 | 
 922 |                 def create_custom_forward(module):
 923 |                     def custom_forward(*inputs):
 924 |                         # None for past_key_value
 925 |                         return module(*inputs, use_cache, output_attentions)
 926 | 
 927 |                     return custom_forward
 928 | 
 929 |                 outputs = torch.utils.checkpoint.checkpoint(
 930 |                     create_custom_forward(block),
 931 |                     hidden_states,
 932 |                     None,
 933 |                     attention_mask,
 934 |                     head_mask[i],
 935 |                     encoder_hidden_states,
 936 |                     encoder_attention_mask,
 937 |                 )
 938 |             else:
 939 |                 outputs = block(
 940 |                     hidden_states,
 941 |                     layer_past=layer_past,
 942 |                     attention_mask=attention_mask,
 943 |                     head_mask=head_mask[i],
 944 |                     encoder_hidden_states=encoder_hidden_states,
 945 |                     encoder_attention_mask=encoder_attention_mask,
 946 |                     use_cache=use_cache,
 947 |                     output_attentions=output_attentions,
 948 |                 )
 949 | 
 950 |             hidden_states = outputs[0]
 951 |             if use_cache is True:
 952 |                 presents = presents + (outputs[1],)
 953 | 
 954 |             if output_attentions:
 955 |                 all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
 956 |                 if self.config.add_cross_attention:
 957 |                     all_cross_attentions = all_cross_attentions + (outputs[3 if use_cache else 2],)
 958 | 
 959 |             # Model Parallel: If it's the last layer for that device, put things on the next device
 960 |             if self.model_parallel:
 961 |                 for k, v in self.device_map.items():
 962 |                     if i == v[-1] and "cuda:" + str(k) != self.last_device:
 963 |                         hidden_states = hidden_states.to("cuda:" + str(k + 1))
 964 | 
 965 |         hidden_states = self.ln_f(hidden_states)
 966 | 
 967 |         hidden_states = hidden_states.view(output_shape)
 968 |         # Add last hidden state
 969 |         if output_hidden_states:
 970 |             all_hidden_states = all_hidden_states + (hidden_states,)
 971 | 
 972 |         if not return_dict:
 973 |             return tuple(
 974 |                 v
 975 |                 for v in [hidden_states, presents, all_hidden_states, all_self_attentions, all_cross_attentions]
 976 |                 if v is not None
 977 |             )
 978 | 
 979 |         return BaseModelOutputWithPastAndCrossAttentions(
 980 |             last_hidden_state=hidden_states,
 981 |             past_key_values=presents,
 982 |             hidden_states=all_hidden_states,
 983 |             attentions=all_self_attentions,
 984 |             cross_attentions=all_cross_attentions,
 985 |         )
 986 | 
 987 | 
 988 | @add_start_docstrings(
 989 |     """
 990 |     The GPT2 Model transformer with a language modeling head on top (linear layer with weights tied to the input
 991 |     embeddings).
 992 |     """,
 993 |     GPT2_START_DOCSTRING,
 994 | )
 995 | class GPT2LMHeadModel(GPT2PreTrainedModel):
 996 |     _keys_to_ignore_on_load_missing = [r"attn.masked_bias", r"attn.bias", r"lm_head.weight"]
 997 | 
 998 |     def __init__(self, config):
 999 |         super().__init__(config)
1000 |         self.transformer = GPT2Model(config)
1001 |         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
1002 | 
1003 |         # Model parallel
1004 |         self.model_parallel = False
1005 |         self.device_map = None
1006 | 
1007 |         # Initialize weights and apply final processing
1008 |         self.post_init()
1009 | 
1010 |     @add_start_docstrings(PARALLELIZE_DOCSTRING)
1011 |     def parallelize(self, device_map=None):
1012 |         self.device_map = (
1013 |             get_device_map(len(self.transformer.h), range(torch.cuda.device_count()))
1014 |             if device_map is None
1015 |             else device_map
1016 |         )
1017 |         assert_device_map(self.device_map, len(self.transformer.h))
1018 |         self.transformer.parallelize(self.device_map)
1019 |         self.lm_head = self.lm_head.to(self.transformer.first_device)
1020 |         self.model_parallel = True
1021 | 
1022 |     @add_start_docstrings(DEPARALLELIZE_DOCSTRING)
1023 |     def deparallelize(self):
1024 |         self.transformer.deparallelize()
1025 |         self.transformer = self.transformer.to("cpu")
1026 |         self.lm_head = self.lm_head.to("cpu")
1027 |         self.model_parallel = False
1028 |         torch.cuda.empty_cache()
1029 | 
1030 |     def get_output_embeddings(self):
1031 |         return self.lm_head
1032 | 
1033 |     def set_output_embeddings(self, new_embeddings):
1034 |         self.lm_head = new_embeddings
1035 | 
1036 |     def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
1037 |         token_type_ids = kwargs.get("token_type_ids", None)
1038 |         # only last token for inputs_ids if past is defined in kwargs
1039 |         if past:
1040 |             input_ids = input_ids[:, -1].unsqueeze(-1)
1041 |             if token_type_ids is not None:
1042 |                 token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
1043 | 
1044 |         attention_mask = kwargs.get("attention_mask", None)
1045 |         position_ids = kwargs.get("position_ids", None)
1046 | 
1047 |         if attention_mask is not None and position_ids is None:
1048 |             # create position_ids on the fly for batch generation
1049 |             position_ids = attention_mask.long().cumsum(-1) - 1
1050 |             position_ids.masked_fill_(attention_mask == 0, 1)
1051 |             if past:
1052 |                 position_ids = position_ids[:, -1].unsqueeze(-1)
1053 |         else:
1054 |             position_ids = None
1055 |         return {
1056 |             "input_ids": input_ids,
1057 |             "past_key_values": past,
1058 |             "use_cache": kwargs.get("use_cache"),
1059 |             "position_ids": position_ids,
1060 |             "attention_mask": attention_mask,
1061 |             "token_type_ids": token_type_ids,
1062 |         }
1063 | 
1064 |     @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
1065 |     @add_code_sample_docstrings(
1066 |         processor_class=_TOKENIZER_FOR_DOC,
1067 |         checkpoint=_CHECKPOINT_FOR_DOC,
1068 |         output_type=CausalLMOutputWithCrossAttentions,
1069 |         config_class=_CONFIG_FOR_DOC,
1070 |     )
1071 |     def forward(
1072 |         self,
1073 |         input_ids=None,
1074 |         past_key_values=None,
1075 |         attention_mask=None,
1076 |         token_type_ids=None,
1077 |         position_ids=None,
1078 |         head_mask=None,
1079 |         inputs_embeds=None,
1080 |         encoder_hidden_states=None,
1081 |         encoder_attention_mask=None,
1082 |         labels=None,
1083 |         use_cache=None,
1084 |         output_attentions=None,
1085 |         output_hidden_states=None,
1086 |         return_dict=None,
1087 |     ):
1088 |         r"""
1089 |         labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1090 |             Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
1091 |             `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
1092 |             are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
1093 |         """
1094 |         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1095 | 
1096 |         transformer_outputs = self.transformer(
1097 |             input_ids,
1098 |             past_key_values=past_key_values,
1099 |             attention_mask=attention_mask,
1100 |             token_type_ids=token_type_ids,
1101 |             position_ids=position_ids,
1102 |             head_mask=head_mask,
1103 |             inputs_embeds=inputs_embeds,
1104 |             encoder_hidden_states=encoder_hidden_states,
1105 |             encoder_attention_mask=encoder_attention_mask,
1106 |             use_cache=use_cache,
1107 |             output_attentions=output_attentions,
1108 |             output_hidden_states=output_hidden_states,
1109 |             return_dict=return_dict,
1110 |         )
1111 |         hidden_states = transformer_outputs[0]
1112 | 
1113 |         # Set device for model parallelism
1114 |         if self.model_parallel:
1115 |             torch.cuda.set_device(self.transformer.first_device)
1116 |             hidden_states = hidden_states.to(self.lm_head.weight.device)
1117 | 
1118 |         lm_logits = self.lm_head(hidden_states)
1119 | 
1120 |         loss = None
1121 |         if labels is not None:
1122 |             # Shift so that tokens < n predict n
1123 |             shift_logits = lm_logits[..., :-1, :].contiguous()
1124 |             shift_labels = labels[..., 1:].contiguous()
1125 |             # Flatten the tokens
1126 |             loss_fct = CrossEntropyLoss()
1127 |             loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
1128 | 
1129 |         if not return_dict:
1130 |             output = (lm_logits,) + transformer_outputs[1:]
1131 |             return ((loss,) + output) if loss is not None else output
1132 | 
1133 |         return CausalLMOutputWithCrossAttentions(
1134 |             loss=loss,
1135 |             logits=lm_logits,
1136 |             past_key_values=transformer_outputs.past_key_values,
1137 |             hidden_states=transformer_outputs.hidden_states,
1138 |             attentions=transformer_outputs.attentions,
1139 |             cross_attentions=transformer_outputs.cross_attentions,
1140 |         )
1141 | 
1142 |     @staticmethod
1143 |     def _reorder_cache(past: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor) -> Tuple[Tuple[torch.Tensor]]:
1144 |         """
1145 |         This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
1146 |         [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
1147 |         beam_idx at every generation step.
1148 |         """
1149 |         return tuple(
1150 |             tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
1151 |             for layer_past in past
1152 |         )
1153 | 
1154 | 
1155 | class BartAttention(nn.Module):
1156 |     """Multi-headed attention from 'Attention Is All You Need' paper"""
1157 | 
1158 |     def __init__(
1159 |         self,
1160 |         embed_dim: int,
1161 |         num_heads: int,
1162 |         dropout: float = 0.0,
1163 |         is_decoder: bool = False,
1164 |         bias: bool = True,
1165 |     ):
1166 |         super().__init__()
1167 |         self.embed_dim = embed_dim
1168 |         self.num_heads = num_heads
1169 |         self.dropout = dropout
1170 |         self.head_dim = embed_dim // num_heads
1171 | 
1172 |         if (self.head_dim * num_heads) != self.embed_dim:
1173 |             raise ValueError(
1174 |                 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
1175 |                 f" and `num_heads`: {num_heads})."
1176 |             )
1177 |         self.scaling = self.head_dim**-0.5
1178 |         self.is_decoder = is_decoder
1179 | 
1180 |         self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
1181 |         self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
1182 |         self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
1183 |         self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
1184 | 
1185 |     def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
1186 |         return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
1187 | 
1188 |     def forward(
1189 |         self,
1190 |         hidden_states: torch.Tensor,
1191 |         key_value_states: Optional[torch.Tensor] = None,
1192 |         past_key_value: Optional[Tuple[torch.Tensor]] = None,
1193 |         attention_mask: Optional[torch.Tensor] = None,
1194 |         layer_head_mask: Optional[torch.Tensor] = None,
1195 |         output_attentions: bool = False,
1196 |     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
1197 |         """Input shape: Batch x Time x Channel"""
1198 | 
1199 |         # if key_value_states are provided this layer is used as a cross-attention layer
1200 |         # for the decoder
1201 |         is_cross_attention = key_value_states is not None
1202 | 
1203 |         bsz, tgt_len, _ = hidden_states.size()
1204 | 
1205 |         # get query proj
1206 |         query_states = self.q_proj(hidden_states) * self.scaling
1207 |         # get key, value proj
1208 |         if is_cross_attention and past_key_value is not None:
1209 |             # reuse k,v, cross_attentions
1210 |             key_states = past_key_value[0]
1211 |             value_states = past_key_value[1]
1212 |         elif is_cross_attention:
1213 |             # cross_attentions
1214 |             key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
1215 |             value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
1216 |         elif past_key_value is not None:
1217 |             # reuse k, v, self_attention
1218 |             key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
1219 |             value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
1220 |             key_states = torch.cat([past_key_value[0], key_states], dim=2)
1221 |             value_states = torch.cat([past_key_value[1], value_states], dim=2)
1222 |         else:
1223 |             # self_attention
1224 |             key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
1225 |             value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
1226 | 
1227 |         if self.is_decoder:
1228 |             # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
1229 |             # Further calls to cross_attention layer can then reuse all cross-attention
1230 |             # key/value_states (first "if" case)
1231 |             # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
1232 |             # all previous decoder key/value_states. Further calls to uni-directional self-attention
1233 |             # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
1234 |             # if encoder bi-directional self-attention `past_key_value` is always `None`
1235 |             past_key_value = (key_states, value_states)
1236 | 
1237 |         proj_shape = (bsz * self.num_heads, -1, self.head_dim)
1238 |         query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
1239 |         key_states = key_states.view(*proj_shape)
1240 |         value_states = value_states.view(*proj_shape)
1241 | 
1242 |         src_len = key_states.size(1)
1243 |         attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
1244 | 
1245 |         if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
1246 |             raise ValueError(
1247 |                 f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
1248 |                 f" {attn_weights.size()}"
1249 |             )
1250 | 
1251 |         if attention_mask is not None:
1252 |             if attention_mask.size() != (bsz, 1, tgt_len, src_len):
1253 |                 raise ValueError(
1254 |                     f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
1255 |                 )
1256 |             attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
1257 |             attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
1258 | 
1259 |         attn_weights = nn.functional.softmax(attn_weights, dim=-1)
1260 | 
1261 |         if layer_head_mask is not None:
1262 |             if layer_head_mask.size() != (self.num_heads,):
1263 |                 raise ValueError(
1264 |                     f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"
1265 |                     f" {layer_head_mask.size()}"
1266 |                 )
1267 |             attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
1268 |             attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
1269 | 
1270 |         if output_attentions:
1271 |             # this operation is a bit awkward, but it's required to
1272 |             # make sure that attn_weights keeps its gradient.
1273 |             # In order to do so, attn_weights have to be reshaped
1274 |             # twice and have to be reused in the following
1275 |             attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
1276 |             attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
1277 |         else:
1278 |             attn_weights_reshaped = None
1279 | 
1280 |         attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
1281 | 
1282 |         attn_output = torch.bmm(attn_probs, value_states)
1283 | 
1284 |         if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
1285 |             raise ValueError(
1286 |                 f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
1287 |                 f" {attn_output.size()}"
1288 |             )
1289 | 
1290 |         attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
1291 |         attn_output = attn_output.transpose(1, 2)
1292 | 
1293 |         # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be
1294 |         # partitioned aross GPUs when using tensor-parallelism.
1295 |         attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
1296 | 
1297 |         attn_output = self.out_proj(attn_output)
1298 | 
1299 |         return attn_output, attn_weights_reshaped, past_key_value
1300 | 
1301 | 
1302 | class SlotAttentionPooling(nn.Module):
1303 |     def __init__(
1304 |         self,
1305 |         input_dim,
1306 |         n_slots,
1307 |     ):
1308 |         super(SlotAttentionPooling, self).__init__()
1309 |         std = input_dim ** -0.5
1310 |         self.W = nn.Parameter(torch.empty(n_slots, input_dim))
1311 |         nn.init.uniform_(self.W, -std, std)
1312 |         self.b = nn.Parameter(torch.empty(n_slots))
1313 |         nn.init.uniform_(self.b, -std, std)
1314 |         self.softmax = nn.functional.softmax
1315 |     
1316 |     def forward(
1317 |         self,
1318 |         batch_rep,
1319 |         attention_mask,
1320 |     ):
1321 |         bsz, n_style, n_slots, input_dim = batch_rep.shape
1322 |         att_logits = (batch_rep * self.W).sum(-1) + self.b
1323 |         att_logits = att_logits.transpose(1, 2)
1324 |         if attention_mask is not None:
1325 |             att_logits = att_logits + ((1.0 - attention_mask) * -10000.0).unsqueeze(1)
1326 |         att_logits = att_logits.reshape(bsz * n_slots, n_style)
1327 |         att_w = self.softmax(att_logits, dim=-1).unsqueeze(-1)
1328 |         batch_rep = batch_rep.transpose(1, 2).reshape(bsz * n_slots, n_style, input_dim)
1329 |         utter_rep = torch.sum(batch_rep * att_w, dim=1).reshape(bsz, n_slots, input_dim)
1330 | 
1331 |         return utter_rep
1332 | 
1333 | 
1334 | class SlotLinear(nn.Module):
1335 |     def __init__(
1336 |         self,
1337 |         input_dim,
1338 |         n_slots,
1339 |     ):
1340 |         super(SlotLinear, self).__init__()
1341 |         std = input_dim ** -0.5
1342 |         self.W = nn.Parameter(torch.empty(n_slots, input_dim, input_dim))
1343 |         nn.init.uniform_(self.W, -std, std)
1344 |         self.b = nn.Parameter(torch.empty(n_slots, input_dim))
1345 |         nn.init.uniform_(self.b, -std, std)
1346 |     
1347 |     def forward(
1348 |         self,
1349 |         feat,
1350 |     ):
1351 |         return torch.einsum('bsi,sij->bsj', [feat, self.W]) + self.b
1352 | 
1353 | 
1354 | # Copied from transformers.models.bart.modeling_bart._expand_mask
1355 | def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
1356 |     """
1357 |     Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
1358 |     """
1359 |     bsz, src_len = mask.size()
1360 |     tgt_len = tgt_len if tgt_len is not None else src_len
1361 | 
1362 |     expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
1363 | 
1364 |     inverted_mask = 1.0 - expanded_mask
1365 | 
1366 |     return inverted_mask.masked_fill(inverted_mask.bool(), torch.finfo(dtype).min)
1367 | 
1368 | 
1369 | class PrefixGPT2(GPT2PreTrainedModel):
1370 |     _keys_to_ignore_on_load_missing = [r"attn.masked_bias", r"attn.bias", r"lm_head.weight"]
1371 | 
1372 |     def __init__(self, config):
1373 |         super().__init__(config)
1374 | 
1375 |         self.n_layer = config.n_layer
1376 |         self.n_head = config.n_head
1377 |         self.d_head = config.n_embd // config.n_head
1378 |         self.n_slots = 30
1379 | 
1380 |         self.dropout = torch.nn.Dropout(config.resid_pdrop)
1381 |         self.style_encoder = AutoModel.from_pretrained("roberta-base", local_files_only=False)
1382 |         self.slot_embedding = nn.Parameter(torch.Tensor(self.n_slots, self.style_encoder.config.hidden_size))
1383 |         nn.init.normal_(self.slot_embedding, mean=0., std=self.style_encoder.config.hidden_size ** -0.5)
1384 |         self.slot_attention =  BartAttention(
1385 |             embed_dim=self.style_encoder.config.hidden_size,
1386 |             num_heads = 1,
1387 |             # num_heads=self.style_encoder.config.num_attention_heads,
1388 |             dropout=config.attn_pdrop,
1389 |         )
1390 |         self.slot_attn_layer_norm = nn.LayerNorm(self.style_encoder.config.hidden_size)
1391 |         self.slot_pooler = SlotAttentionPooling(self.style_encoder.config.hidden_size, self.n_slots)
1392 |         self.slot_pooler_layer_norm = nn.LayerNorm(self.style_encoder.config.hidden_size)
1393 |         self.slot_proj = SlotLinear(self.style_encoder.config.hidden_size, self.n_slots)
1394 |         self.prefix_proj = torch.nn.Sequential(
1395 |             torch.nn.Linear(self.style_encoder.config.hidden_size, self.style_encoder.config.hidden_size),
1396 |             torch.nn.Tanh(),
1397 |             torch.nn.Linear(self.style_encoder.config.hidden_size, config.n_embd * 2 * config.n_layer)
1398 |         )
1399 | 
1400 |         self.transformer = GPT2Model(config)
1401 |         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
1402 | 
1403 |         # Model parallel
1404 |         self.model_parallel = False
1405 |         self.device_map = None
1406 | 
1407 |         # Initialize weights and apply final processing
1408 |         self.post_init()
1409 | 
1410 |     # @add_start_docstrings(PARALLELIZE_DOCSTRING)
1411 |     def parallelize(self, device_map=None):
1412 |         self.device_map = (
1413 |             get_device_map(len(self.transformer.h), range(torch.cuda.device_count()))
1414 |             if device_map is None
1415 |             else device_map
1416 |         )
1417 |         assert_device_map(self.device_map, len(self.transformer.h))
1418 |         self.transformer.parallelize(self.device_map)
1419 |         self.lm_head = self.lm_head.to(self.transformer.first_device)
1420 |         self.model_parallel = True
1421 | 
1422 |     # @add_start_docstrings(DEPARALLELIZE_DOCSTRING)
1423 |     def deparallelize(self):
1424 |         self.transformer.deparallelize()
1425 |         self.transformer = self.transformer.to("cpu")
1426 |         self.lm_head = self.lm_head.to("cpu")
1427 |         self.model_parallel = False
1428 |         torch.cuda.empty_cache()
1429 | 
1430 |     def get_output_embeddings(self):
1431 |         return self.lm_head
1432 | 
1433 |     def set_output_embeddings(self, new_embeddings):
1434 |         self.lm_head = new_embeddings
1435 | 
1436 |     def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
1437 |         token_type_ids = kwargs.get("token_type_ids", None)
1438 |         # only last token for inputs_ids if past is defined in kwargs
1439 |         if past:
1440 |             input_ids = input_ids[:, -1].unsqueeze(-1)
1441 |             if token_type_ids is not None:
1442 |                 token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
1443 | 
1444 |         attention_mask = kwargs.get("attention_mask", None)
1445 |         position_ids = kwargs.get("position_ids", None)
1446 | 
1447 |         if attention_mask is not None and position_ids is None:
1448 |             # create position_ids on the fly for batch generation
1449 |             position_ids = attention_mask.long().cumsum(-1) - 1
1450 |             position_ids.masked_fill_(attention_mask == 0, 1)
1451 |             if past:
1452 |                 position_ids = position_ids[:, -1].unsqueeze(-1)
1453 |         else:
1454 |             position_ids = None
1455 |         return {
1456 |             "input_ids": input_ids,
1457 |             "ref_input_ids": kwargs.get("ref_input_ids"),
1458 |             "past_key_values": past,
1459 |             "use_cache": kwargs.get("use_cache"),
1460 |             "position_ids": position_ids,
1461 |             "attention_mask": attention_mask,
1462 |             "ref_attention_mask": kwargs.get("ref_attention_mask"),
1463 |             "token_type_ids": token_type_ids,
1464 |         }
1465 | 
1466 |     def get_prompt(self, hidden_states):
1467 |         past_key_values = hidden_states.view(
1468 |             hidden_states.size(0),
1469 |             self.n_slots,
1470 |             self.n_layer * 2, 
1471 |             self.n_head,
1472 |             self.d_head
1473 |         )
1474 |         past_key_values = self.dropout(past_key_values)
1475 |         past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2)
1476 |         return past_key_values
1477 | 
1478 |     def pool_style(self, h_style, num_style, attention_mask):
1479 |         bsz_mul_num_style, style_len, d_model = h_style.shape
1480 |         bsz = bsz_mul_num_style // num_style
1481 |         utter_mask = attention_mask.any(dim=-1).to(attention_mask.dtype)
1482 |         style_mask = utter_mask.reshape(bsz, -1).any(dim=-1).to(utter_mask.dtype).unsqueeze(1).repeat(1, self.n_slots)
1483 |         utter_slot_reps = self.slot_attention(
1484 |             hidden_states=self.slot_embedding.unsqueeze(0).expand(bsz_mul_num_style, self.n_slots, d_model),
1485 |             key_value_states=h_style,
1486 |             attention_mask=_expand_mask(
1487 |                 attention_mask.reshape(bsz_mul_num_style, style_len), 
1488 |                 h_style.dtype, 
1489 |                 tgt_len=self.n_slots,
1490 |             ),
1491 |         )[0]
1492 |         utter_slot_reps = self.slot_attn_layer_norm(utter_slot_reps).reshape(bsz, num_style, self.n_slots, d_model)
1493 |         style_slot_reps = self.slot_pooler_layer_norm(self.slot_pooler(utter_slot_reps, utter_mask))
1494 | 
1495 |         return style_slot_reps, style_mask
1496 | 
1497 |     # @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
1498 |     # @add_code_sample_docstrings(
1499 |     #     processor_class=_TOKENIZER_FOR_DOC,
1500 |     #     checkpoint=_CHECKPOINT_FOR_DOC,
1501 |     #     output_type=CausalLMOutputWithCrossAttentions,
1502 |     #     config_class=_CONFIG_FOR_DOC,
1503 |     # )
1504 |     def forward(
1505 |         self,
1506 |         input_ids=None,
1507 |         ref_input_ids=None,
1508 |         past_key_values=None,
1509 |         attention_mask=None,
1510 |         ref_attention_mask=None,
1511 |         token_type_ids=None,
1512 |         position_ids=None,
1513 |         head_mask=None,
1514 |         inputs_embeds=None,
1515 |         encoder_hidden_states=None,
1516 |         encoder_attention_mask=None,
1517 |         labels=None,
1518 |         use_cache=None,
1519 |         output_attentions=None,
1520 |         output_hidden_states=None,
1521 |         return_dict=None,
1522 |     ):
1523 |         r"""
1524 |         labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1525 |             Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
1526 |             `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
1527 |             are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
1528 |         """
1529 |         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1530 | 
1531 |         if past_key_values is None:
1532 |             bsz = input_ids.size(0)
1533 |             style_bsz, max_n_style, max_style_len = ref_input_ids.shape
1534 |             assert bsz == style_bsz
1535 | 
1536 |             style_outputs = self.style_encoder(
1537 |                 input_ids=ref_input_ids.reshape(-1, max_style_len),
1538 |                 attention_mask=ref_attention_mask.reshape(-1, max_style_len),
1539 |                 head_mask=head_mask,
1540 |                 inputs_embeds=inputs_embeds,
1541 |                 output_attentions=output_attentions,
1542 |                 output_hidden_states=output_hidden_states,
1543 |                 return_dict=return_dict,
1544 |             )
1545 |             style_pool_hidden, style_mask = self.pool_style(style_outputs["last_hidden_state"], max_n_style, ref_attention_mask)
1546 |             style_pool_hidden = self.slot_proj(style_pool_hidden)
1547 |             style_proj_hidden = self.prefix_proj(style_pool_hidden)
1548 |             past_key_values = self.get_prompt(style_proj_hidden)
1549 |         else:
1550 |             assert not self.training
1551 | 
1552 |         prefix_attention_mask = torch.ones(input_ids.size(0), self.n_slots).to(self.transformer.device)
1553 |         attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
1554 | 
1555 |         if position_ids is not None:
1556 |             position_ids += self.n_slots
1557 | 
1558 |         transformer_outputs = self.transformer(
1559 |             input_ids,
1560 |             past_key_values=past_key_values,
1561 |             attention_mask=attention_mask,
1562 |             token_type_ids=token_type_ids,
1563 |             position_ids=position_ids,
1564 |             head_mask=head_mask,
1565 |             inputs_embeds=inputs_embeds,
1566 |             encoder_hidden_states=encoder_hidden_states,
1567 |             encoder_attention_mask=encoder_attention_mask,
1568 |             use_cache=use_cache,
1569 |             output_attentions=output_attentions,
1570 |             output_hidden_states=output_hidden_states,
1571 |             return_dict=return_dict,
1572 |         )
1573 |         hidden_states = transformer_outputs[0]
1574 | 
1575 |         # Set device for model parallelism
1576 |         if self.model_parallel:
1577 |             torch.cuda.set_device(self.transformer.first_device)
1578 |             hidden_states = hidden_states.to(self.lm_head.weight.device)
1579 | 
1580 |         lm_logits = self.lm_head(hidden_states)
1581 | 
1582 |         loss = None
1583 |         if labels is not None:
1584 |             # Shift so that tokens < n predict n
1585 |             shift_logits = lm_logits[..., :-1, :].contiguous()
1586 |             shift_labels = labels[..., 1:].contiguous()
1587 |             # Flatten the tokens
1588 |             loss_fct = CrossEntropyLoss()
1589 |             loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
1590 | 
1591 |         if not return_dict:
1592 |             output = (lm_logits,) + transformer_outputs[1:]
1593 |             return ((loss,) + output) if loss is not None else output
1594 | 
1595 |         return CausalLMOutputWithCrossAttentions(
1596 |             loss=loss,
1597 |             logits=lm_logits,
1598 |             past_key_values=transformer_outputs.past_key_values,
1599 |             hidden_states=transformer_outputs.hidden_states,
1600 |             attentions=transformer_outputs.attentions,
1601 |             cross_attentions=transformer_outputs.cross_attentions,
1602 |         )
1603 | 
1604 |     @staticmethod
1605 |     def _reorder_cache(past: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor) -> Tuple[Tuple[torch.Tensor]]:
1606 |         """
1607 |         This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
1608 |         [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
1609 |         beam_idx at every generation step.
1610 |         """
1611 |         return tuple(
1612 |             tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
1613 |             for layer_past in past
1614 |         )
1615 | 
1616 |     def estimate_tokens(self, input_dict: Dict[str, Union[torch.Tensor, Any]]) -> int:
1617 |         """
1618 |         Helper function to estimate the total number of tokens from the model inputs.
1619 |         Args:
1620 |             inputs (`dict`): The model inputs.
1621 |         Returns:
1622 |             `int`: The total number of tokens.
1623 |         """
1624 |         if self.main_input_name in input_dict:
1625 |             return input_dict[self.main_input_name][0].numel() + input_dict[self.main_input_name][1].numel()
1626 |         else:
1627 |             logger.warning(
1628 |                 "Could not estimate the number of tokens of the input, floating-point operations will not be computed"
1629 |             )
1630 |             return 0
1631 | 
1632 | 
1633 | class ContextPrefixGPT2(GPT2PreTrainedModel):
1634 |     _keys_to_ignore_on_load_missing = [r"attn.masked_bias", r"attn.bias", r"lm_head.weight"]
1635 | 
1636 |     def __init__(self, config):
1637 |         super().__init__(config)
1638 | 
1639 |         self.n_layer = config.n_layer
1640 |         self.n_head = config.n_head
1641 |         self.d_head = config.n_embd // config.n_head
1642 |         self.n_slots = 30
1643 | 
1644 |         self.dropout = torch.nn.Dropout(config.resid_pdrop)
1645 |         self.style_encoder = AutoModel.from_pretrained("roberta-base", local_files_only=True)
1646 |         self.slot_embedding = nn.Parameter(torch.Tensor(self.n_slots, self.style_encoder.config.hidden_size))
1647 |         nn.init.normal_(self.slot_embedding, mean=0., std=self.style_encoder.config.hidden_size ** -0.5)
1648 |         self.context_attention = BartAttention(
1649 |             embed_dim=self.style_encoder.config.hidden_size,
1650 |             num_heads = 1,
1651 |             # num_heads=self.style_encoder.config.num_attention_heads,
1652 |             dropout=config.attn_pdrop,
1653 |         )
1654 |         self.context_attn_layer_norm = nn.LayerNorm(self.style_encoder.config.hidden_size)
1655 | 
1656 |         # self.ref_pos_embed = nn.Parameter(torch.Tensor(100, self.style_encoder.config.hidden_size))
1657 |         # nn.init.normal_(self.ref_pos_embed, mean=0., std=self.style_encoder.config.hidden_size ** -0.5)
1658 |         self.gatew = nn.Linear(self.style_encoder.config.hidden_size * 2, 1)
1659 | 
1660 |         self.ref_attention = BartAttention(
1661 |             embed_dim=self.style_encoder.config.hidden_size,
1662 |             num_heads = 1,
1663 |             # num_heads=self.style_encoder.config.num_attention_heads,
1664 |             dropout=config.attn_pdrop,
1665 |         )
1666 |         self.ref_attn_layer_norm = nn.LayerNorm(self.style_encoder.config.hidden_size)
1667 |         self.prefix_proj = torch.nn.Sequential(
1668 |             torch.nn.Linear(self.style_encoder.config.hidden_size, self.style_encoder.config.hidden_size),
1669 |             torch.nn.Tanh(),
1670 |             torch.nn.Linear(self.style_encoder.config.hidden_size, config.n_embd * 2 * config.n_layer)
1671 |         )
1672 | 
1673 |         self.transformer = GPT2Model(config)
1674 |         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
1675 | 
1676 |         # Model parallel
1677 |         self.model_parallel = False
1678 |         self.device_map = None
1679 | 
1680 |         # Initialize weights and apply final processing
1681 |         self.post_init()
1682 | 
1683 |     # @add_start_docstrings(PARALLELIZE_DOCSTRING)
1684 |     def parallelize(self, device_map=None):
1685 |         self.device_map = (
1686 |             get_device_map(len(self.transformer.h), range(torch.cuda.device_count()))
1687 |             if device_map is None
1688 |             else device_map
1689 |         )
1690 |         assert_device_map(self.device_map, len(self.transformer.h))
1691 |         self.transformer.parallelize(self.device_map)
1692 |         self.lm_head = self.lm_head.to(self.transformer.first_device)
1693 |         self.model_parallel = True
1694 | 
1695 |     # @add_start_docstrings(DEPARALLELIZE_DOCSTRING)
1696 |     def deparallelize(self):
1697 |         self.transformer.deparallelize()
1698 |         self.transformer = self.transformer.to("cpu")
1699 |         self.lm_head = self.lm_head.to("cpu")
1700 |         self.model_parallel = False
1701 |         torch.cuda.empty_cache()
1702 | 
1703 |     def get_output_embeddings(self):
1704 |         return self.lm_head
1705 | 
1706 |     def set_output_embeddings(self, new_embeddings):
1707 |         self.lm_head = new_embeddings
1708 | 
1709 |     def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
1710 |         token_type_ids = kwargs.get("token_type_ids", None)
1711 |         # only last token for inputs_ids if past is defined in kwargs
1712 |         if past:
1713 |             input_ids = input_ids[:, -1].unsqueeze(-1)
1714 |             if token_type_ids is not None:
1715 |                 token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
1716 | 
1717 |         attention_mask = kwargs.get("attention_mask", None)
1718 |         position_ids = kwargs.get("position_ids", None)
1719 | 
1720 |         if attention_mask is not None and position_ids is None:
1721 |             # create position_ids on the fly for batch generation
1722 |             position_ids = attention_mask.long().cumsum(-1) - 1
1723 |             position_ids.masked_fill_(attention_mask == 0, 1)
1724 |             if past:
1725 |                 position_ids = position_ids[:, -1].unsqueeze(-1)
1726 |         else:
1727 |             position_ids = None
1728 |         return {
1729 |             "input_ids": input_ids,
1730 |             "context_input_ids": kwargs.get("context_input_ids"),
1731 |             "ref_input_ids": kwargs.get("ref_input_ids"),
1732 |             "past_key_values": past,
1733 |             "use_cache": kwargs.get("use_cache"),
1734 |             "position_ids": position_ids,
1735 |             "attention_mask": attention_mask,
1736 |             "context_attention_mask": kwargs.get("context_attention_mask"),
1737 |             "ref_attention_mask": kwargs.get("ref_attention_mask"),
1738 |             "token_type_ids": token_type_ids,
1739 |         }
1740 | 
1741 |     def get_prompt(self, hidden_states):
1742 |         past_key_values = hidden_states.view(
1743 |             hidden_states.size(0),
1744 |             self.n_slots,
1745 |             self.n_layer * 2, 
1746 |             self.n_head,
1747 |             self.d_head
1748 |         )
1749 |         past_key_values = self.dropout(past_key_values)
1750 |         past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2)
1751 |         return past_key_values
1752 | 
1753 |     def add_ref_pos(self, h_ref, num_ref):
1754 |         """
1755 |         h_ref.shape=(bsz * num_ref, ref_len, d_model)
1756 |         """
1757 |         # # static ref positional embedding
1758 |         # bsz_mul_num_ref, ref_len, d_model = h_ref.shape
1759 |         # bsz = bsz_mul_num_ref // num_ref
1760 |         # # (bsz, num_ref, ref_len, d_model)
1761 |         # h_ref = h_ref.reshape(bsz, num_ref, ref_len, d_model)
1762 |         # # (bsz, num_ref, ref_len, d_model)
1763 |         # embed_pos = self.ref_pos_embed[:num_ref].unsqueeze(0).unsqueeze(2).expand(bsz, num_ref, ref_len, d_model)
1764 |         # h_ref = h_ref + embed_pos
1765 |         # h_ref = h_ref.reshape(bsz * num_ref, ref_len, d_model)
1766 | 
1767 |         # dynamic ref position embedding
1768 |         bsz_mul_num_ref, ref_len, d_model = h_ref.shape
1769 |         bsz = bsz_mul_num_ref // num_ref
1770 |         # (bsz, num_ref, ref_len, d_model)
1771 |         h_ref = h_ref.reshape(bsz, num_ref, ref_len, d_model)
1772 |         # (bsz, num_ref, ref_len, d_model)
1773 |         utter = h_ref[:, :, 0, :].unsqueeze(-2).expand(bsz, num_ref, ref_len, d_model)
1774 |         # (bsz, num_ref, ref_len, 1)
1775 |         coef = self.gatew(torch.cat([h_ref, utter], dim=-1)).sigmoid()
1776 |         # (bsz, num_ref, ref_len, d_model)
1777 |         h_ref = coef * h_ref + (1 - coef) * utter
1778 |         h_ref = h_ref.reshape(bsz * num_ref, ref_len, d_model)
1779 |         
1780 |         return h_ref
1781 | 
1782 |     # @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
1783 |     # @add_code_sample_docstrings(
1784 |     #     processor_class=_TOKENIZER_FOR_DOC,
1785 |     #     checkpoint=_CHECKPOINT_FOR_DOC,
1786 |     #     output_type=CausalLMOutputWithCrossAttentions,
1787 |     #     config_class=_CONFIG_FOR_DOC,
1788 |     # )
1789 |     def forward(
1790 |         self,
1791 |         input_ids=None,
1792 |         context_input_ids=None,
1793 |         ref_input_ids=None,
1794 |         past_key_values=None,
1795 |         attention_mask=None,
1796 |         context_attention_mask=None,
1797 |         ref_attention_mask=None,
1798 |         token_type_ids=None,
1799 |         position_ids=None,
1800 |         head_mask=None,
1801 |         inputs_embeds=None,
1802 |         encoder_hidden_states=None,
1803 |         encoder_attention_mask=None,
1804 |         labels=None,
1805 |         use_cache=None,
1806 |         output_attentions=None,
1807 |         output_hidden_states=None,
1808 |         return_dict=None,
1809 |     ):
1810 |         r"""
1811 |         labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1812 |             Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
1813 |             `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
1814 |             are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
1815 |         """
1816 |         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1817 | 
1818 |         if past_key_values is None:
1819 |             bsz = input_ids.size(0)
1820 |             style_bsz, max_n_style, max_style_len = ref_input_ids.shape
1821 |             assert bsz == style_bsz
1822 | 
1823 |             context_outputs = self.style_encoder(
1824 |                 input_ids=context_input_ids,
1825 |                 attention_mask=context_attention_mask,
1826 |                 head_mask=None,
1827 |                 inputs_embeds=None,
1828 |                 output_attentions=output_attentions,
1829 |                 output_hidden_states=output_hidden_states,
1830 |                 return_dict=return_dict,
1831 |             )
1832 |             projected_context = self.context_attention(
1833 |                 hidden_states=self.slot_embedding.unsqueeze(0).expand(bsz, self.n_slots, self.style_encoder.config.hidden_size),
1834 |                 key_value_states=context_outputs["last_hidden_state"],
1835 |                 attention_mask=_expand_mask(
1836 |                     context_attention_mask, 
1837 |                     context_outputs["last_hidden_state"].dtype, 
1838 |                     tgt_len=self.n_slots,
1839 |                 ),
1840 |             )[0]
1841 |             projected_context = self.context_attn_layer_norm(projected_context)
1842 |             style_outputs = self.style_encoder(
1843 |                 input_ids=ref_input_ids.reshape(-1, max_style_len),
1844 |                 attention_mask=ref_attention_mask.reshape(-1, max_style_len),
1845 |                 head_mask=None,
1846 |                 inputs_embeds=None,
1847 |                 output_attentions=output_attentions,
1848 |                 output_hidden_states=output_hidden_states,
1849 |                 return_dict=return_dict,
1850 |             )
1851 |             style_hidden_states = self.add_ref_pos(style_outputs["last_hidden_state"], max_n_style).reshape(style_bsz, max_n_style * max_style_len, -1)
1852 |             # style_hidden_states = style_outputs["last_hidden_state"].reshape(style_bsz, max_n_style * max_style_len, -1)
1853 |             style_attention_mask = ref_attention_mask.reshape(style_bsz, -1)
1854 |             projected_style = self.ref_attention(
1855 |                 hidden_states=projected_context,
1856 |                 key_value_states=style_hidden_states,
1857 |                 attention_mask=_expand_mask(
1858 |                     style_attention_mask, 
1859 |                     style_hidden_states.dtype, 
1860 |                     tgt_len=self.n_slots,
1861 |                 ),
1862 |             )[0]
1863 |             projected_style = self.ref_attn_layer_norm(projected_style)
1864 |             projected_style = self.prefix_proj(projected_style)
1865 |             past_key_values = self.get_prompt(projected_style)
1866 |         else:
1867 |             assert not self.training
1868 | 
1869 |         if position_ids is not None:
1870 |             position_ids += self.n_slots
1871 | 
1872 |         prefix_attention_mask = torch.ones(input_ids.size(0), self.n_slots).to(self.transformer.device)
1873 |         attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
1874 | 
1875 |         transformer_outputs = self.transformer(
1876 |             input_ids,
1877 |             past_key_values=past_key_values,
1878 |             attention_mask=attention_mask,
1879 |             token_type_ids=token_type_ids,
1880 |             position_ids=position_ids,
1881 |             head_mask=head_mask,
1882 |             inputs_embeds=inputs_embeds,
1883 |             encoder_hidden_states=encoder_hidden_states,
1884 |             encoder_attention_mask=encoder_attention_mask,
1885 |             use_cache=use_cache,
1886 |             output_attentions=output_attentions,
1887 |             output_hidden_states=output_hidden_states,
1888 |             return_dict=return_dict,
1889 |         )
1890 |         hidden_states = transformer_outputs[0]
1891 | 
1892 |         # Set device for model parallelism
1893 |         if self.model_parallel:
1894 |             torch.cuda.set_device(self.transformer.first_device)
1895 |             hidden_states = hidden_states.to(self.lm_head.weight.device)
1896 | 
1897 |         lm_logits = self.lm_head(hidden_states)
1898 | 
1899 |         loss = None
1900 |         if labels is not None:
1901 |             # Shift so that tokens < n predict n
1902 |             shift_logits = lm_logits[..., :-1, :].contiguous()
1903 |             shift_labels = labels[..., 1:].contiguous()
1904 |             # Flatten the tokens
1905 |             loss_fct = CrossEntropyLoss()
1906 |             loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
1907 | 
1908 |         if not return_dict:
1909 |             output = (lm_logits,) + transformer_outputs[1:]
1910 |             return ((loss,) + output) if loss is not None else output
1911 | 
1912 |         return CausalLMOutputWithCrossAttentions(
1913 |             loss=loss,
1914 |             logits=lm_logits,
1915 |             past_key_values=transformer_outputs.past_key_values,
1916 |             hidden_states=transformer_outputs.hidden_states,
1917 |             attentions=transformer_outputs.attentions,
1918 |             cross_attentions=transformer_outputs.cross_attentions,
1919 |         )
1920 | 
1921 |     @staticmethod
1922 |     def _reorder_cache(past: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor) -> Tuple[Tuple[torch.Tensor]]:
1923 |         """
1924 |         This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
1925 |         [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
1926 |         beam_idx at every generation step.
1927 |         """
1928 |         return tuple(
1929 |             tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
1930 |             for layer_past in past
1931 |         )
1932 | 
1933 |     def estimate_tokens(self, input_dict: Dict[str, Union[torch.Tensor, Any]]) -> int:
1934 |         """
1935 |         Helper function to estimate the total number of tokens from the model inputs.
1936 |         Args:
1937 |             inputs (`dict`): The model inputs.
1938 |         Returns:
1939 |             `int`: The total number of tokens.
1940 |         """
1941 |         if self.main_input_name in input_dict:
1942 |             return input_dict[self.main_input_name][0].numel() + input_dict[self.main_input_name][1].numel()
1943 |         else:
1944 |             logger.warning(
1945 |                 "Could not estimate the number of tokens of the input, floating-point operations will not be computed"
1946 |             )
1947 |             return 0
1948 | 
1949 | 
1950 | class DHAP(GPT2PreTrainedModel):
1951 |     _keys_to_ignore_on_load_missing = [r"attn.masked_bias", r"attn.bias", r"lm_head.weight"]
1952 | 
1953 |     def __init__(self, config):
1954 |         super().__init__(config)
1955 |         tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small", use_fast=True, local_files_only=True)
1956 |         tokenizer.pad_token = tokenizer.eos_token
1957 |         roberta_tokenizer = AutoTokenizer.from_pretrained("distilroberta-base", use_fast=True, local_files_only=True)
1958 |         token_id_map = [tokenizer.convert_tokens_to_ids(roberta_tokenizer.convert_ids_to_tokens(i)) for i in range(50265)]
1959 |         self.token_id_map = torch.tensor(token_id_map)
1960 | 
1961 |         self.n_layer = config.n_layer
1962 |         self.n_head = config.n_head
1963 |         self.d_head = config.n_embd // config.n_head
1964 |         self.n_slots = 30
1965 | 
1966 |         ### dhap components
1967 |         self.hist_resp_encoder = AutoModel.from_pretrained("distilroberta-base")
1968 |         self.post_encoder = self.hist_post_encoder = AutoModel.from_pretrained("distilroberta-base")
1969 |         self.transformer = GPT2Model(config)
1970 |         self.post_output_attn = nn.MultiheadAttention(768, 1, batch_first=True)
1971 |         self.hist_attn = nn.MultiheadAttention(768, 1, batch_first=True)
1972 |         self.copy_attn = nn.MultiheadAttention(768, 1, batch_first=True)
1973 |         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
1974 |         self.personal_lm_head = nn.Linear(config.n_embd * 3, config.vocab_size, bias=False)
1975 |         self.switcher = nn.Linear(config.n_embd * 4, 2)
1976 | 
1977 |         ### prefix components
1978 |         self.dropout = torch.nn.Dropout(config.resid_pdrop)
1979 |         self.slot_embedding = nn.Parameter(torch.Tensor(self.n_slots, self.hist_resp_encoder.config.hidden_size))
1980 |         nn.init.normal_(self.slot_embedding, mean=0., std=self.hist_resp_encoder.config.hidden_size ** -0.5)
1981 |         self.ctx_proj_attn = nn.MultiheadAttention(768, 1, batch_first=True, dropout=config.attn_pdrop)
1982 |         self.ctx_proj_attn_layer_norm = nn.LayerNorm(self.hist_resp_encoder.config.hidden_size)
1983 |         self.gatew = nn.Linear(self.hist_resp_encoder.config.hidden_size * 2, 1)
1984 |         self.ref_proj_attn = nn.MultiheadAttention(768, 1, batch_first=True, dropout=config.attn_pdrop)
1985 |         self.ref_proj_attn_layer_norm = nn.LayerNorm(self.hist_resp_encoder.config.hidden_size)
1986 |         self.prefix_proj = torch.nn.Sequential(
1987 |             torch.nn.Linear(self.hist_resp_encoder.config.hidden_size, self.hist_resp_encoder.config.hidden_size),
1988 |             torch.nn.Tanh(),
1989 |             torch.nn.Linear(self.hist_resp_encoder.config.hidden_size, config.n_embd * 2 * config.n_layer)
1990 |         )
1991 | 
1992 |         # Model parallel
1993 |         self.model_parallel = False
1994 |         self.device_map = None
1995 | 
1996 |         # Initialize weights and apply final processing
1997 |         self.post_init()
1998 | 
1999 |     def parallelize(self, device_map=None):
2000 |         self.device_map = (
2001 |             get_device_map(len(self.transformer.h), range(torch.cuda.device_count()))
2002 |             if device_map is None
2003 |             else device_map
2004 |         )
2005 |         assert_device_map(self.device_map, len(self.transformer.h))
2006 |         self.transformer.parallelize(self.device_map)
2007 |         self.lm_head = self.lm_head.to(self.transformer.first_device)
2008 |         self.model_parallel = True
2009 | 
2010 |     def deparallelize(self):
2011 |         self.transformer.deparallelize()
2012 |         self.transformer = self.transformer.to("cpu")
2013 |         self.lm_head = self.lm_head.to("cpu")
2014 |         self.model_parallel = False
2015 |         torch.cuda.empty_cache()
2016 | 
2017 |     def get_output_embeddings(self):
2018 |         return self.lm_head
2019 | 
2020 |     def set_output_embeddings(self, new_embeddings):
2021 |         self.lm_head = new_embeddings
2022 | 
2023 |     def pre_encode(
2024 |         self,
2025 |         context_ids,
2026 |         hist_p_ids,
2027 |         hist_r_ids,
2028 |         context_attention_mask=None,
2029 |         hist_p_attention_mask=None,
2030 |         hist_r_attention_mask=None,
2031 |     ):
2032 |         bsz = hist_p_ids.size(0)
2033 |         hist_num = hist_p_ids.size(1)
2034 |         hist_p_len = hist_p_ids.size(-1)
2035 |         hist_r_len = hist_r_ids.size(-1)
2036 | 
2037 |         encoded_post = self.post_encoder(input_ids=context_ids, attention_mask=context_attention_mask).last_hidden_state
2038 | 
2039 |         projected_post, _ = self.ctx_proj_attn(
2040 |             self.slot_embedding.unsqueeze(0).expand(bsz, self.n_slots, 768),
2041 |             encoded_post,
2042 |             encoded_post,
2043 |             key_padding_mask=~context_attention_mask.bool(),
2044 |         )
2045 |         projected_post = self.ctx_proj_attn_layer_norm(projected_post)
2046 | 
2047 |         encoded_hist_p = self.hist_post_encoder(input_ids=hist_p_ids.reshape(bsz*hist_num, -1), attention_mask=hist_p_attention_mask.reshape(bsz*hist_num, -1)).last_hidden_state
2048 |         encoded_hist_p = encoded_hist_p.reshape(bsz, hist_num, hist_p_len, -1)
2049 |         hist_p_attention_mask[:, 0, 0] = 0
2050 |         hist_p_attention_mask[hist_p_ids == 2] = 0
2051 |         encoded_hist_p[~hist_p_attention_mask.bool()] = 0
2052 |         encoded_hist_p = encoded_hist_p.mean(2)
2053 | 
2054 |         hist_r_attention_mask[hist_r_ids == 0] = 0
2055 |         hist_r_ids[hist_r_ids == 0] = 1
2056 |         hist_r_attention_mask[:, 0, 0] = 1
2057 |         hist_r_ids[:, 0, 0] = 0
2058 |         encoded_hist_r = self.hist_resp_encoder(input_ids=hist_r_ids.reshape(bsz, -1), attention_mask=hist_r_attention_mask.reshape(bsz, -1)).last_hidden_state
2059 |         general_profile = torch.clone(encoded_hist_r[:, 0])
2060 | 
2061 |         projected_hist, _ = self.ref_proj_attn(
2062 |             projected_post,
2063 |             encoded_hist_r,
2064 |             encoded_hist_r,
2065 |             key_padding_mask=~hist_r_attention_mask.reshape(bsz, -1).bool(),
2066 |         )
2067 |         projected_hist = self.ref_proj_attn_layer_norm(projected_hist)
2068 |         ref_prefix = self.prefix_proj(projected_hist)
2069 |         past_key_values = self.get_prompt(ref_prefix)
2070 | 
2071 |         encoded_hist_r = encoded_hist_r.reshape(bsz, hist_num, hist_r_len, -1)
2072 |         hist_r_attention_mask[:, 0, 0] = 0
2073 |         hist_r_attention_mask[hist_r_ids == 2] = 0
2074 |         encoded_hist_r[~hist_r_attention_mask.bool()] = 0
2075 |         encoded_hist_r_tokens = encoded_hist_r.reshape(bsz, hist_num*hist_r_len, -1)
2076 |         encoded_hist_r = encoded_hist_r.mean(2)
2077 | 
2078 |         return {
2079 |             "encoded_post": encoded_post,
2080 |             "encoded_hist_p": encoded_hist_p,
2081 |             "encoded_hist_r": encoded_hist_r,
2082 |             "encoded_hist_r_tokens": encoded_hist_r_tokens,
2083 |             "general_profile": general_profile,
2084 |             "context_attention_mask": context_attention_mask,
2085 |             "hist_p_attention_mask": hist_p_attention_mask,
2086 |             "hist_r_attention_mask": hist_r_attention_mask,
2087 |             "past_key_values": past_key_values,
2088 |         }
2089 | 
2090 | 
2091 |     def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
2092 |         token_type_ids = kwargs.get("token_type_ids", None)
2093 |         # only last token for inputs_ids if past is defined in kwargs
2094 |         if past:
2095 |             input_ids = input_ids[:, -1].unsqueeze(-1)
2096 |             if token_type_ids is not None:
2097 |                 token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
2098 | 
2099 |         attention_mask = kwargs.get("attention_mask", None)
2100 |         position_ids = kwargs.get("position_ids", None)
2101 | 
2102 |         if attention_mask is not None and position_ids is None:
2103 |             # create position_ids on the fly for batch generation
2104 |             position_ids = attention_mask.long().cumsum(-1) - 1
2105 |             position_ids.masked_fill_(attention_mask == 0, 1)
2106 |             if past:
2107 |                 position_ids = position_ids[:, -1].unsqueeze(-1)
2108 |         else:
2109 |             position_ids = None
2110 | 
2111 |         pre_encoded = self.pre_encode(
2112 |             kwargs.get("context_ids"),
2113 |             kwargs.get("hist_p_ids"),
2114 |             kwargs.get("hist_r_ids"),
2115 |             kwargs.get("context_attention_mask"),
2116 |             kwargs.get("hist_p_attention_mask"),
2117 |             kwargs.get("hist_r_attention_mask"),
2118 |         )
2119 | 
2120 |         return {
2121 |             "input_ids": input_ids,
2122 |             "context_ids": kwargs.get("context_ids"),
2123 |             "hist_p_ids": kwargs.get("hist_p_ids"),
2124 |             "hist_r_ids": kwargs.get("hist_r_ids"),
2125 |             "past_key_values": past,
2126 |             "use_cache": kwargs.get("use_cache"),
2127 |             "position_ids": position_ids,
2128 |             "attention_mask": attention_mask,
2129 |             "context_attention_mask": kwargs.get("context_attention_mask"),
2130 |             "hist_p_attention_mask": kwargs.get("hist_p_attention_mask"),
2131 |             "hist_r_attention_mask": kwargs.get("hist_r_attention_mask"),
2132 |             "token_type_ids": token_type_ids,
2133 |             "pre_encoded": pre_encoded,
2134 |         }
2135 | 
2136 |     def get_prompt(self, hidden_states):
2137 |         past_key_values = hidden_states.view(
2138 |             hidden_states.size(0),
2139 |             self.n_slots,
2140 |             self.n_layer * 2, 
2141 |             self.n_head,
2142 |             self.d_head
2143 |         )
2144 |         past_key_values = self.dropout(past_key_values)
2145 |         past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2)
2146 |         return past_key_values
2147 | 
2148 |     def forward(
2149 |         self,
2150 |         input_ids=None,
2151 |         context_ids=None,
2152 |         hist_p_ids=None,
2153 |         hist_r_ids=None,
2154 |         past_key_values=None,
2155 |         attention_mask=None,
2156 |         context_attention_mask=None,
2157 |         hist_p_attention_mask=None,
2158 |         hist_r_attention_mask=None,
2159 |         pre_encoded=None,
2160 |         generate=False,
2161 |         token_type_ids=None,
2162 |         position_ids=None,
2163 |         head_mask=None,
2164 |         inputs_embeds=None,
2165 |         encoder_hidden_states=None,
2166 |         encoder_attention_mask=None,
2167 |         labels=None,
2168 |         use_cache=None,
2169 |         output_attentions=None,
2170 |         output_hidden_states=None,
2171 |         return_dict=None,
2172 |     ):
2173 |         r"""
2174 |         labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
2175 |             Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
2176 |             `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
2177 |             are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
2178 |         """
2179 |         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
2180 | 
2181 |         bsz = hist_p_ids.size(0)
2182 | 
2183 |         if pre_encoded is None:
2184 |             pre_encoded = self.pre_encode(
2185 |                 context_ids,
2186 |                 hist_p_ids,
2187 |                 hist_r_ids,
2188 |                 context_attention_mask,
2189 |                 hist_p_attention_mask,
2190 |                 hist_r_attention_mask,
2191 |             )
2192 |         encoded_post = pre_encoded["encoded_post"]
2193 |         encoded_hist_p = pre_encoded["encoded_hist_p"]
2194 |         encoded_hist_r = pre_encoded["encoded_hist_r"]
2195 |         encoded_hist_r_tokens = pre_encoded["encoded_hist_r_tokens"]
2196 |         general_profile = pre_encoded["general_profile"]
2197 |         context_attention_mask = pre_encoded["context_attention_mask"]
2198 |         hist_p_attention_mask = pre_encoded["hist_p_attention_mask"]
2199 |         hist_r_attention_mask = pre_encoded["hist_r_attention_mask"]
2200 |         # if past_key_values is None:
2201 |         #     past_key_values = pre_encoded["past_key_values"]
2202 | 
2203 |         # if position_ids is not None:
2204 |         #     position_ids += self.n_slots
2205 |         # prefix_attention_mask = torch.ones(input_ids.size(0), self.n_slots).to(self.transformer.device)
2206 |         # attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
2207 | 
2208 |         transformer_outputs = self.transformer(
2209 |             input_ids,
2210 |             past_key_values=past_key_values,
2211 |             attention_mask=attention_mask,
2212 |             token_type_ids=token_type_ids,
2213 |             position_ids=position_ids,
2214 |             head_mask=head_mask,
2215 |             inputs_embeds=inputs_embeds,
2216 |             encoder_hidden_states=encoder_hidden_states,
2217 |             encoder_attention_mask=encoder_attention_mask,
2218 |             use_cache=use_cache,
2219 |             output_attentions=output_attentions,
2220 |             output_hidden_states=output_hidden_states,
2221 |             return_dict=return_dict,
2222 |         )
2223 | 
2224 |         transformer_hidden_states = transformer_outputs[0]
2225 |         expanded_general_profile = general_profile.unsqueeze(1).expand(-1, transformer_hidden_states.size(1), -1)
2226 |         dyn_post, _ = self.post_output_attn(transformer_hidden_states, encoded_post, encoded_post, key_padding_mask=~context_attention_mask.bool())
2227 |         dyn_know, _ = self.hist_attn(transformer_hidden_states, encoded_hist_p, encoded_hist_r)
2228 |         _, copy_weights = self.copy_attn(dyn_post, encoded_hist_r_tokens, encoded_hist_r_tokens, key_padding_mask=~hist_r_attention_mask.bool().reshape(bsz, -1))
2229 |         mapped_hist_r_ids = self.token_id_map[hist_r_ids].reshape(bsz, -1).to(hist_r_ids.device)
2230 | 
2231 |         logits = self.lm_head(transformer_hidden_states)
2232 |         logits += self.personal_lm_head(torch.cat([expanded_general_profile, dyn_post, dyn_know], dim=2))
2233 |         lm_logits = logits
2234 |         probs = nn.functional.softmax(logits, -1)
2235 |         switch = self.switcher(torch.cat([transformer_hidden_states, expanded_general_profile, dyn_post, dyn_know], dim=2))
2236 |         switch = nn.functional.softmax(switch, -1)
2237 |         copy_probs = torch.zeros_like(probs).cuda()
2238 |         copy_probs = copy_probs.scatter_add(2, mapped_hist_r_ids.unsqueeze(1).expand(-1, logits.size(1), -1), copy_weights)
2239 |         combined_probs = probs * switch[:, :, 0].unsqueeze(-1) + copy_probs * switch[:, :, 1].unsqueeze(-1)
2240 |         combined_probs[combined_probs == 0] = 1e-6
2241 |         lm_logits = combined_probs.log()
2242 |         lm_logits = lm_logits - lm_logits.min(-1)[0].unsqueeze(-1)
2243 | 
2244 |         loss = None
2245 |         if labels is not None:
2246 |             # Shift so that tokens < n predict n
2247 |             shift_logits = lm_logits[..., :-1, :].contiguous()
2248 |             shift_labels = labels[..., 1:].contiguous()
2249 |             # Flatten the tokens
2250 |             loss_fct = CrossEntropyLoss()
2251 |             loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
2252 | 
2253 |             # penalize copying
2254 |             # loss += switch[:, :, 1].mean()
2255 | 
2256 |         if not return_dict:
2257 |             output = (lm_logits,) + transformer_outputs[1:]
2258 |             return ((loss,) + output) if loss is not None else output
2259 | 
2260 |         return CausalLMOutputWithCrossAttentions(
2261 |             loss=loss,
2262 |             logits=lm_logits,
2263 |             past_key_values=transformer_outputs.past_key_values,
2264 |             hidden_states=transformer_outputs.hidden_states,
2265 |             attentions=transformer_outputs.attentions,
2266 |             cross_attentions=transformer_outputs.cross_attentions,
2267 |         )
2268 | 
2269 |     @staticmethod
2270 |     def _reorder_cache(past: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor) -> Tuple[Tuple[torch.Tensor]]:
2271 |         """
2272 |         This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
2273 |         [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
2274 |         beam_idx at every generation step.
2275 |         """
2276 |         return tuple(
2277 |             tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
2278 |             for layer_past in past
2279 |         )
2280 |         
2281 | 
2282 | @add_start_docstrings(
2283 |     """
2284 | The GPT2 Model transformer with a language modeling and a multiple-choice classification head on top e.g. for
2285 | RocStories/SWAG tasks. The two heads are two linear layers. The language modeling head has its weights tied to the
2286 | input embeddings, the classification head takes as input the input of a specified classification token index in the
2287 | input sequence).
2288 | """,
2289 |     GPT2_START_DOCSTRING,
2290 | )
2291 | class GPT2DoubleHeadsModel(GPT2PreTrainedModel):
2292 |     _keys_to_ignore_on_load_missing = [r"attn.masked_bias", r"attn.bias", r"lm_head.weight"]
2293 | 
2294 |     def __init__(self, config):
2295 |         super().__init__(config)
2296 |         config.num_labels = 1
2297 |         self.transformer = GPT2Model(config)
2298 |         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
2299 |         self.multiple_choice_head = SequenceSummary(config)
2300 | 
2301 |         # Model parallel
2302 |         self.model_parallel = False
2303 |         self.device_map = None
2304 | 
2305 |         # Initialize weights and apply final processing
2306 |         self.post_init()
2307 | 
2308 |     @add_start_docstrings(PARALLELIZE_DOCSTRING)
2309 |     def parallelize(self, device_map=None):
2310 |         self.device_map = (
2311 |             get_device_map(len(self.transformer.h), range(torch.cuda.device_count()))
2312 |             if device_map is None
2313 |             else device_map
2314 |         )
2315 |         assert_device_map(self.device_map, len(self.transformer.h))
2316 |         self.transformer.parallelize(self.device_map)
2317 |         self.lm_head = self.lm_head.to(self.transformer.first_device)
2318 |         self.multiple_choice_head = self.multiple_choice_head.to(self.transformer.first_device)
2319 |         self.model_parallel = True
2320 | 
2321 |     @add_start_docstrings(DEPARALLELIZE_DOCSTRING)
2322 |     def deparallelize(self):
2323 |         self.transformer.deparallelize()
2324 |         self.transformer = self.transformer.to("cpu")
2325 |         self.lm_head = self.lm_head.to("cpu")
2326 |         self.multiple_choice_head = self.multiple_choice_head.to("cpu")
2327 |         self.model_parallel = False
2328 |         torch.cuda.empty_cache()
2329 | 
2330 |     def get_output_embeddings(self):
2331 |         return self.lm_head
2332 | 
2333 |     def set_output_embeddings(self, new_embeddings):
2334 |         self.lm_head = new_embeddings
2335 | 
2336 |     def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
2337 |         token_type_ids = kwargs.get("token_type_ids", None)
2338 |         # only last token for inputs_ids if past is defined in kwargs
2339 |         if past:
2340 |             input_ids = input_ids[:, -1].unsqueeze(-1)
2341 |             if token_type_ids is not None:
2342 |                 token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
2343 | 
2344 |         attention_mask = kwargs.get("attention_mask", None)
2345 |         position_ids = kwargs.get("position_ids", None)
2346 | 
2347 |         if attention_mask is not None and position_ids is None:
2348 |             # create position_ids on the fly for batch generation
2349 |             position_ids = attention_mask.long().cumsum(-1) - 1
2350 |             position_ids.masked_fill_(attention_mask == 0, 1)
2351 |             if past:
2352 |                 position_ids = position_ids[:, -1].unsqueeze(-1)
2353 |         else:
2354 |             position_ids = None
2355 | 
2356 |         return {
2357 |             "input_ids": input_ids,
2358 |             "past_key_values": past,
2359 |             "use_cache": kwargs.get("use_cache"),
2360 |             "position_ids": position_ids,
2361 |             "attention_mask": attention_mask,
2362 |             "token_type_ids": token_type_ids,
2363 |         }
2364 | 
2365 |     @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
2366 |     @replace_return_docstrings(output_type=GPT2DoubleHeadsModelOutput, config_class=_CONFIG_FOR_DOC)
2367 |     def forward(
2368 |         self,
2369 |         input_ids=None,
2370 |         past_key_values=None,
2371 |         attention_mask=None,
2372 |         token_type_ids=None,
2373 |         position_ids=None,
2374 |         head_mask=None,
2375 |         inputs_embeds=None,
2376 |         mc_token_ids=None,
2377 |         labels=None,
2378 |         mc_labels=None,
2379 |         use_cache=None,
2380 |         output_attentions=None,
2381 |         output_hidden_states=None,
2382 |         return_dict=None,
2383 |         **kwargs,
2384 |     ):
2385 |         r"""
2386 |         mc_token_ids (`torch.LongTensor` of shape `(batch_size, num_choices)`, *optional*, default to index of the last token of the input):
2387 |             Index of the classification token in each input sequence. Selected in the range `[0, input_ids.size(-1) -
2388 |             1[`.
2389 |         labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
2390 |             Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
2391 |             `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size - 1]` All labels set to
2392 |             `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size - 1]`
2393 |         mc_labels (`torch.LongTensor` of shape `(batch_size)`, *optional*):
2394 |             Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., num_choices]`
2395 |             where *num_choices* is the size of the second dimension of the input tensors. (see *input_ids* above)
2396 | 
2397 |         Return:
2398 | 
2399 |         Example:
2400 | 
2401 |         ```python
2402 |         >>> import torch
2403 |         >>> from transformers import GPT2Tokenizer, GPT2DoubleHeadsModel
2404 | 
2405 |         >>> tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
2406 |         >>> model = GPT2DoubleHeadsModel.from_pretrained("gpt2")
2407 | 
2408 |         >>> # Add a [CLS] to the vocabulary (we should train it also!)
2409 |         >>> num_added_tokens = tokenizer.add_special_tokens({"cls_token": "[CLS]"})
2410 | 
2411 |         >>> embedding_layer = model.resize_token_embeddings(
2412 |         ...     len(tokenizer)
2413 |         >>> )  # Update the model embeddings with the new vocabulary size
2414 | 
2415 |         >>> choices = ["Hello, my dog is cute [CLS]", "Hello, my cat is cute [CLS]"]
2416 |         >>> encoded_choices = [tokenizer.encode(s) for s in choices]
2417 |         >>> cls_token_location = [tokens.index(tokenizer.cls_token_id) for tokens in encoded_choices]
2418 | 
2419 |         >>> input_ids = torch.tensor(encoded_choices).unsqueeze(0)  # Batch size: 1, number of choices: 2
2420 |         >>> mc_token_ids = torch.tensor([cls_token_location])  # Batch size: 1
2421 | 
2422 |         >>> outputs = model(input_ids, mc_token_ids=mc_token_ids)
2423 |         >>> lm_logits = outputs.logits
2424 |         >>> mc_logits = outputs.mc_logits
2425 |         ```"""
2426 |         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
2427 | 
2428 |         transformer_outputs = self.transformer(
2429 |             input_ids,
2430 |             past_key_values=past_key_values,
2431 |             attention_mask=attention_mask,
2432 |             token_type_ids=token_type_ids,
2433 |             position_ids=position_ids,
2434 |             head_mask=head_mask,
2435 |             inputs_embeds=inputs_embeds,
2436 |             use_cache=use_cache,
2437 |             output_attentions=output_attentions,
2438 |             output_hidden_states=output_hidden_states,
2439 |             return_dict=return_dict,
2440 |         )
2441 | 
2442 |         hidden_states = transformer_outputs[0]
2443 | 
2444 |         # Set device for model parallelism
2445 |         if self.model_parallel:
2446 |             torch.cuda.set_device(self.transformer.first_device)
2447 |             hidden_states = hidden_states.to(self.lm_head.weight.device)
2448 | 
2449 |         lm_logits = self.lm_head(hidden_states)
2450 |         mc_logits = self.multiple_choice_head(hidden_states, mc_token_ids).squeeze(-1)
2451 | 
2452 |         mc_loss = None
2453 |         if mc_labels is not None:
2454 |             loss_fct = CrossEntropyLoss()
2455 |             mc_loss = loss_fct(mc_logits.view(-1, mc_logits.size(-1)), mc_labels.view(-1))
2456 |         lm_loss = None
2457 |         if labels is not None:
2458 |             shift_logits = lm_logits[..., :-1, :].contiguous()
2459 |             shift_labels = labels[..., 1:].contiguous()
2460 |             loss_fct = CrossEntropyLoss()
2461 |             lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
2462 | 
2463 |         if not return_dict:
2464 |             output = (lm_logits, mc_logits) + transformer_outputs[1:]
2465 |             if mc_loss is not None:
2466 |                 output = (mc_loss,) + output
2467 |             return ((lm_loss,) + output) if lm_loss is not None else output
2468 | 
2469 |         return GPT2DoubleHeadsModelOutput(
2470 |             loss=lm_loss,
2471 |             mc_loss=mc_loss,
2472 |             logits=lm_logits,
2473 |             mc_logits=mc_logits,
2474 |             past_key_values=transformer_outputs.past_key_values,
2475 |             hidden_states=transformer_outputs.hidden_states,
2476 |             attentions=transformer_outputs.attentions,
2477 |         )
2478 | 
2479 |     @staticmethod
2480 |     def _reorder_cache(past: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor) -> Tuple[Tuple[torch.Tensor]]:
2481 |         """
2482 |         This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
2483 |         [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
2484 |         beam_idx at every generation step.
2485 |         """
2486 |         return tuple(
2487 |             tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
2488 |             for layer_past in past
2489 |         )
2490 | 
2491 | 
2492 | @add_start_docstrings(
2493 |     """
2494 |     The GPT2 Model transformer with a sequence classification head on top (linear layer).
2495 | 
2496 |     [`GPT2ForSequenceClassification`] uses the last token in order to do the classification, as other causal models
2497 |     (e.g. GPT-1) do.
2498 | 
2499 |     Since it does classification on the last token, it requires to know the position of the last token. If a
2500 |     `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
2501 |     no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
2502 |     padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
2503 |     each row of the batch).
2504 |     """,
2505 |     GPT2_START_DOCSTRING,
2506 | )
2507 | class GPT2ForSequenceClassification(GPT2PreTrainedModel):
2508 |     _keys_to_ignore_on_load_missing = [r"h\.\d+\.attn\.masked_bias", r"lm_head\.weight"]
2509 | 
2510 |     def __init__(self, config):
2511 |         super().__init__(config)
2512 |         self.num_labels = config.num_labels
2513 |         self.transformer = GPT2Model(config)
2514 |         self.score = nn.Linear(config.n_embd, self.num_labels, bias=False)
2515 | 
2516 |         # Model parallel
2517 |         self.model_parallel = False
2518 |         self.device_map = None
2519 | 
2520 |         # Initialize weights and apply final processing
2521 |         self.post_init()
2522 | 
2523 |     @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
2524 |     @add_code_sample_docstrings(
2525 |         processor_class=_TOKENIZER_FOR_DOC,
2526 |         checkpoint="microsoft/DialogRPT-updown",
2527 |         output_type=SequenceClassifierOutputWithPast,
2528 |         config_class=_CONFIG_FOR_DOC,
2529 |     )
2530 |     def forward(
2531 |         self,
2532 |         input_ids=None,
2533 |         past_key_values=None,
2534 |         attention_mask=None,
2535 |         token_type_ids=None,
2536 |         position_ids=None,
2537 |         head_mask=None,
2538 |         inputs_embeds=None,
2539 |         labels=None,
2540 |         use_cache=None,
2541 |         output_attentions=None,
2542 |         output_hidden_states=None,
2543 |         return_dict=None,
2544 |     ):
2545 |         r"""
2546 |         labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
2547 |             Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
2548 |             config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
2549 |             `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
2550 |         """
2551 |         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
2552 | 
2553 |         transformer_outputs = self.transformer(
2554 |             input_ids,
2555 |             past_key_values=past_key_values,
2556 |             attention_mask=attention_mask,
2557 |             token_type_ids=token_type_ids,
2558 |             position_ids=position_ids,
2559 |             head_mask=head_mask,
2560 |             inputs_embeds=inputs_embeds,
2561 |             use_cache=use_cache,
2562 |             output_attentions=output_attentions,
2563 |             output_hidden_states=output_hidden_states,
2564 |             return_dict=return_dict,
2565 |         )
2566 |         hidden_states = transformer_outputs[0]
2567 |         logits = self.score(hidden_states)
2568 | 
2569 |         if input_ids is not None:
2570 |             batch_size, sequence_length = input_ids.shape[:2]
2571 |         else:
2572 |             batch_size, sequence_length = inputs_embeds.shape[:2]
2573 | 
2574 |         assert (
2575 |             self.config.pad_token_id is not None or batch_size == 1
2576 |         ), "Cannot handle batch sizes > 1 if no padding token is defined."
2577 |         if self.config.pad_token_id is None:
2578 |             sequence_lengths = -1
2579 |         else:
2580 |             if input_ids is not None:
2581 |                 sequence_lengths = torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1
2582 |             else:
2583 |                 sequence_lengths = -1
2584 |                 logger.warning(
2585 |                     f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
2586 |                     f"unexpected if using padding tokens in conjunction with `inputs_embeds.`"
2587 |                 )
2588 | 
2589 |         pooled_logits = logits[torch.arange(batch_size, device=self.device), sequence_lengths]
2590 | 
2591 |         loss = None
2592 |         if labels is not None:
2593 |             if self.config.problem_type is None:
2594 |                 if self.num_labels == 1:
2595 |                     self.config.problem_type = "regression"
2596 |                 elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
2597 |                     self.config.problem_type = "single_label_classification"
2598 |                 else:
2599 |                     self.config.problem_type = "multi_label_classification"
2600 | 
2601 |             if self.config.problem_type == "regression":
2602 |                 loss_fct = MSELoss()
2603 |                 if self.num_labels == 1:
2604 |                     loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
2605 |                 else:
2606 |                     loss = loss_fct(pooled_logits, labels)
2607 |             elif self.config.problem_type == "single_label_classification":
2608 |                 loss_fct = CrossEntropyLoss()
2609 |                 loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
2610 |             elif self.config.problem_type == "multi_label_classification":
2611 |                 loss_fct = BCEWithLogitsLoss()
2612 |                 loss = loss_fct(pooled_logits, labels)
2613 |         if not return_dict:
2614 |             output = (pooled_logits,) + transformer_outputs[1:]
2615 |             return ((loss,) + output) if loss is not None else output
2616 | 
2617 |         return SequenceClassifierOutputWithPast(
2618 |             loss=loss,
2619 |             logits=pooled_logits,
2620 |             past_key_values=transformer_outputs.past_key_values,
2621 |             hidden_states=transformer_outputs.hidden_states,
2622 |             attentions=transformer_outputs.attentions,
2623 |         )
2624 | 
2625 | 
2626 | @add_start_docstrings(
2627 |     """
2628 |     GPT2 Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
2629 |     Named-Entity-Recognition (NER) tasks.
2630 |     """,
2631 |     GPT2_START_DOCSTRING,
2632 | )
2633 | class GPT2ForTokenClassification(GPT2PreTrainedModel):
2634 |     def __init__(self, config):
2635 |         super().__init__(config)
2636 |         self.num_labels = config.num_labels
2637 | 
2638 |         self.transformer = GPT2Model(config)
2639 |         if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:
2640 |             classifier_dropout = config.classifier_dropout
2641 |         elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None:
2642 |             classifier_dropout = config.hidden_dropout
2643 |         else:
2644 |             classifier_dropout = 0.1
2645 |         self.dropout = nn.Dropout(classifier_dropout)
2646 |         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
2647 | 
2648 |         # Model parallel
2649 |         self.model_parallel = False
2650 |         self.device_map = None
2651 | 
2652 |         # Initialize weights and apply final processing
2653 |         self.post_init()
2654 | 
2655 |     @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
2656 |     @add_code_sample_docstrings(
2657 |         processor_class=_TOKENIZER_FOR_DOC,
2658 |         checkpoint="microsoft/DialogRPT-updown",
2659 |         output_type=TokenClassifierOutput,
2660 |         config_class=_CONFIG_FOR_DOC,
2661 |     )
2662 |     def forward(
2663 |         self,
2664 |         input_ids=None,
2665 |         past_key_values=None,
2666 |         attention_mask=None,
2667 |         token_type_ids=None,
2668 |         position_ids=None,
2669 |         head_mask=None,
2670 |         inputs_embeds=None,
2671 |         labels=None,
2672 |         use_cache=None,
2673 |         output_attentions=None,
2674 |         output_hidden_states=None,
2675 |         return_dict=None,
2676 |     ):
2677 |         r"""
2678 |         labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
2679 |             Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
2680 |             config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
2681 |             `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
2682 |         """
2683 |         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
2684 | 
2685 |         transformer_outputs = self.transformer(
2686 |             input_ids,
2687 |             past_key_values=past_key_values,
2688 |             attention_mask=attention_mask,
2689 |             token_type_ids=token_type_ids,
2690 |             position_ids=position_ids,
2691 |             head_mask=head_mask,
2692 |             inputs_embeds=inputs_embeds,
2693 |             use_cache=use_cache,
2694 |             output_attentions=output_attentions,
2695 |             output_hidden_states=output_hidden_states,
2696 |             return_dict=return_dict,
2697 |         )
2698 | 
2699 |         hidden_states = transformer_outputs[0]
2700 |         hidden_states = self.dropout(hidden_states)
2701 |         logits = self.classifier(hidden_states)
2702 | 
2703 |         loss = None
2704 |         if labels is not None:
2705 |             loss_fct = CrossEntropyLoss()
2706 |             loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
2707 | 
2708 |         if not return_dict:
2709 |             output = (logits,) + transformer_outputs[2:]
2710 |             return ((loss,) + output) if loss is not None else output
2711 | 
2712 |         return TokenClassifierOutput(
2713 |             loss=loss,
2714 |             logits=logits,
2715 |             hidden_states=transformer_outputs.hidden_states,
2716 |             attentions=transformer_outputs.attentions,
2717 |         )


--------------------------------------------------------------------------------