wassname/Open-Assistant
1
0
Fork 0
mirror of https://github.com/wassname/Open-Assistant.git synced 2026-09-09 11:15:08 +08:00
Code Issues Packages Projects Releases Wiki Activity

Question-Answer special tokens

Browse Source
This commit is contained in:
Sotirios Anagnostidis
2023-01-03 22:02:32 +01:00
parent e14c670cb0
commit 3a10e9412d
4 changed files with 38 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model/supervised_finetuning/custom_datasets/__init__.py
+6
View File
@@ -2,6 +2,12 @@ from datasets import load_dataset
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset, Subset
QA_SPECIAL_TOKENS = {
'Question': '<question>',
'Answer': '<answer>'
}
class SquadV2Dataset(Dataset):
def __init__(self, cache_dir, split):
model/supervised_finetuning/custom_datasets/dialogue_collator.py
+12 -9
View File
@@ -6,6 +6,8 @@ import torch
from torch.nn import functional as F
from transformers.tokenization_utils_base import PaddingStrategy, PreTrainedTokenizerBase
from . import QA_SPECIAL_TOKENS
@dataclass
class DialogueDataCollator:
@@ -19,22 +21,21 @@ class DialogueDataCollator:
pad_to_multiple_of: Optional[int] = None
def __call__(self, features):
# TODO add special tokens for question and answer here
# additional_special_tokens = ['<question>', '<answer>']
prompt_tokens = ["Question: ", "Answer: "]
flatten_messages = []
label_masks = []
for messages in features:
assert len(messages) % 2 == 0, "Number of messages must be even"
messages = [
(prompt_tokens[0] if i % 2 == 0 else "") + x + ((" " + prompt_tokens[1]) if i % 2 == 0 else "")
(QA_SPECIAL_TOKENS["Question"] if i % 2 == 0 else "")
+ x
+ (QA_SPECIAL_TOKENS["Answer"] if i % 2 == 0 else "")
for i, x in enumerate(messages)
]
# Add a way for the model to terminate generation, reinitialize prompter
messages.append(prompt_tokens[0])
# Add a way for the model to terminate generation
# When we predict the start of a new expected question, we want to be able to stop generation
messages.append(QA_SPECIAL_TOKENS["Question"])
flatten_messages.append(
self.tokenizer(
@@ -47,8 +48,10 @@ class DialogueDataCollator:
message_change_indices = np.cumsum([len(x) for x in messages[:-1]])
# for each token an integer indicating the index of the message it belongs to. Just to create the label mask.
# TEXT: Question: Hello, how are you? Answer: I am fine. Question: What is your name? Answer: My name is John.
# MESSAGE_INDICES: 0 0 0 0 0 0 1 1 1 2 2 2 2 2 2 3 3 3 3
# Label mask is true when predicting a token that is part of the answer, false otherwise.
# TEXT: Question: Hello, how are you? Answer: I am fine. Question: What is your name? Answer: My name is John. Question:
# MESSAGE_INDICES: 0 0 0 0 0 0 1 1 1 2 2 2 2 2 2 3 3 3 3 -2
# LABEL_MASK: 0 0 0 0 0 1 1 1 1 0 0 0 0 0 1 1 1 1 1 0
# If no result in next, we are predicting the last termination token(s)
message_indices = list(