Upload 5 files

README.md

@@ -1,3 +1,7 @@
----
-license: apache-2.0
----
+## Training procedure
+
+### Framework versions
+
+- PEFT 0.4.0
+
+- PEFT 0.4.0

adapter_config.json

@@ -0,0 +1,26 @@
+{
+  "auto_mapping": null,
+  "base_model_name_or_path": "mistralai/Mistral-7B-Instruct-v0.1",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "lora_alpha": 64,
+  "lora_dropout": 0.05,
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 64,
+  "revision": null,
+  "target_modules": [
+    "q_proj",
+    "k_proj",
+    "v_proj",
+    "o_proj",
+    "gate_proj",
+    "up_proj",
+    "down_proj"
+  ],
+  "task_type": "CAUSAL_LM"
+}

adapter_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:04ea77ff2bb1943c52e1554007240514abd8bf16a9bb1b47e925a27a71bf555a
+size 671250189

inference.py

@@ -0,0 +1,163 @@
+import torch
+torch.cuda.empty_cache()
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+base_model = "mistralai/Mistral-7B-Instruct-v0.1"
+# new_model = "kmichiru/Nikaido-7B-mistral-instruct-v0.1"
+new_model = "kmichiru/Nikaido-7B-mistral-instruct-v0.3-vn_v2"
+
+# Reload tokenizer
+tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
+print(tokenizer.pad_token, tokenizer.pad_token_id)
+tokenizer.padding_side = "right"
+
+# Reload the base model
+base_model_reload = AutoModelForCausalLM.from_pretrained(
+    base_model, low_cpu_mem_usage=True,
+    return_dict=True,torch_dtype=torch.bfloat16,
+    device_map= {"": 0})
+model = PeftModel.from_pretrained(base_model_reload, new_model)
+# model = model.merge_and_unload()
+
+
+
+model.config.use_cache = True
+model.eval()
+
+def dialogue(role, content):
+    return {
+        "role": role,
+        "content": content
+    }
+
+
+import json, random
+TRAIN_DSET = "iroseka_dataset.jsonl"
+try:
+    with open(TRAIN_DSET, "r", encoding="utf-8") as f:
+        examples = [json.loads(line) for line in f]
+except FileNotFoundError:
+    print("Few-shot data not found, skipping...")
+    examples = []
+
+def format_chat_history(example, few_shot=0):
+    user_msgs = []
+    # for inference each round, we only need the user messages
+    for msg in example["messages"]:
+        # if msg["role"] == "user":
+        user_msgs.append(msg["content"]) 
+    messages = [
+        dialogue("user", "\n".join(user_msgs)), # join user messages together
+        # example["messages"][-1], # the last message is the bot's response
+    ]
+
+    if few_shot > 0:
+        # randomly sample a few messages from the dialogue history
+        few_shot_data = random.sample(examples, few_shot)
+        for few_shot_example in few_shot_data:
+            few_shot_msgs = []
+            for msg in few_shot_example["messages"]:
+                if msg["role"] == "user":
+                    few_shot_msgs.append(msg["content"])
+            messages = [
+                dialogue("user", "\n".join(few_shot_msgs)),
+                few_shot_example["messages"][-1]
+            ] + messages
+
+    encodeds = tokenizer.apply_chat_template(messages, tokenize=False)
+    return encodeds
+
+def format_chat_history_v2(example, few_shot):
+    # TODO: implement few-shot learning
+    user_msg = []
+    user_msg.append("<s>")
+    for msg in example["messages"]:
+        # [INST] What is your favourite condiment? [/INST]
+        user_msg.append(f"[INST] {msg['content']} [/INST]")
+    # user_msg.append("</s>")
+    if "next_speaker" in example:
+        user_msg.append(f"[INST] {example['next_speaker']}: ")
+    return " ".join(user_msg)
+
+from transformers import StoppingCriteria, StoppingCriteriaList
+class StoppingCriteriaSub(StoppingCriteria):
+    def __init__(self, stops = [], encounters=1):
+        super().__init__()
+        self.stops = [stop.to("cuda") for stop in stops]
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        for seq in input_ids:
+            for stop in self.stops:
+                if len(seq) >= len(stop) and torch.all((stop == seq[-len(stop):])).item():
+                    return True
+        return False
+    
+stop_words = ["[/INST]"]
+stop_words_ids = [tokenizer(stop_word, return_tensors='pt', add_special_tokens=False)['input_ids'].squeeze() for stop_word in stop_words]
+stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])
+
+def inference(chat_history):
+    # chat_history: dict, with "messages" key storing dialogue history, in OpenAI format
+    formatted = format_chat_history_v2(chat_history, few_shot=1)
+    print(formatted)
+    model_inputs = tokenizer(
+        [formatted],
+        return_tensors="pt",
+    )
+    print(model_inputs)
+    model_inputs = model_inputs.to(model.device)
+    with torch.no_grad():
+        outputs = model.generate(
+            input_ids=model_inputs.input_ids,
+            attention_mask=model_inputs.attention_mask,
+            # max_length=1024,
+            do_sample=True,
+            top_p=1,
+            # contrastive search
+            # top_k=50,
+            # penalty_alpha=0.6,
+            # num_return_sequences=1,
+            temperature=0.3,
+            # num_return_sequences=3,
+            use_cache=True,
+            # pad_token_id=tokenizer.eos_token_id, # eos_token_id is not available for some models
+            pad_token_id=tokenizer.pad_token_id, # eos_token_id is not available for some models
+            eos_token_id=tokenizer.eos_token_id,
+            bos_token_id=tokenizer.bos_token_id,
+            output_scores=True,
+            output_attentions=False,
+            output_hidden_states=False,
+            max_new_tokens=256,
+            # num_beams=9,
+            # num_beam_groups=3,
+            # repetition_penalty=1.0,
+            # diversity_penalty=0.5,
+            # num_beams=5,
+            # stopping_criteria=stopping_criteria,
+        )
+        # print(outputs)
+        text = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        def postprocess(t):
+            t = t.split("[/INST]")
+            t = [x.replace("[INST]", "").strip() for x in t]
+            t = [x for x in t if x != ""]
+            return t[-1]
+        # text = [postprocess(t) for t in text]
+        
+    return text
+
+    
+if __name__ == "__main__":
+    chat_history = {
+        "messages": [
+            # dialogue("system", ""),
+            dialogue("user", "傍白: 真紅の言葉が胸の中に滑り込んでくる。"),
+            dialogue("user", "悠馬: っ"),
+            dialogue("user", "傍白: 限界だった。"),
+            dialogue("user", "悠馬: 真紅，大好きです。これからもずっと一緒にいてください。"),
+        ],
+        "next_speaker": "真紅"
+    }
+    print(inference(chat_history))
+

isft_mistral.py

@@ -0,0 +1,187 @@
+from datasets import load_dataset
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+import os
+
+
+base_model_id = "mistralai/Mistral-7B-Instruct-v0.1" 
+
+WORK = "vn_v2"
+new_model_id = f"kmichiru/Nikaido-7B-mistral-instruct-v0.3-{WORK}"
+
+# DSET = {
+#     "train": f"dataset_iroseka/{WORK}_dataset.jsonl",
+#     "eval": f"dataset_iroseka/{WORK}_validations.jsonl"
+# }
+
+DSET = {
+    "train": f"dataset_iroseka/{WORK}_train.jsonl",
+    "eval": f"dataset_iroseka/{WORK}_val.jsonl"
+}
+
+
+dataset = load_dataset("json", data_files=DSET)
+# model = AutoModelForCausalLM.from_pretrained(base_model_id, torch_dtype=torch.bfloat16)
+tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+# max_length = 1024
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "right"
+
+def dialogue(role, content):
+    return {
+        "role": role,
+        "content": content
+    }
+
+def format_chat_history(example):
+    user_msgs = []
+    for msg in example["messages"]:
+        if msg["role"] == "user":
+            user_msgs.append(msg["content"])
+    messages = [
+        dialogue("user", "\n".join(user_msgs)), # join user messages together
+        example["messages"][-1], # the last message is the bot's response
+    ]
+    encodeds = tokenizer.apply_chat_template(messages, tokenize=False)
+    return encodeds
+
+def prep_speaker(msg: str):
+    msg = msg.replace("\u3000", " ") # replace full-width spaces
+    speaker, content = msg.split(":", 1)
+    speaker = speaker.strip()
+    content = content.strip()
+    if len(speaker) == 0:
+        speaker = "傍白"
+
+    return f"{speaker}: {content}"
+    
+
+def format_chat_history_v2(example):
+    user_msg = []
+    user_msg.append("<s>")
+    for msg in example["messages"]:
+        # [INST] What is your favourite condiment? [/INST]
+        if msg["role"] != "system":
+            user_msg.append(f"[INST] {prep_speaker(msg['content'])} [/INST]")
+    # user_msg.append("</s>")
+    return " ".join(user_msg)
+
+# def format_chat_history_v2(example):
+#     user_msgs = []
+#     for msg in example["messages"]:
+#         if msg["role"] == "user":
+#             user_msgs.append(msg["content"])
+#     messages = [
+#         dialogue("user", "\n".join(user_msgs)), # join user messages together
+#         example["messages"][-1], # the last message is the bot's response
+#     ]
+#     encodeds = tokenizer.apply_chat_template(messages, tokenize=False)
+#     return encodeds
+
+print(format_chat_history_v2(dataset['train'][0]))
+
+def generate_and_tokenize_prompt(prompt, max_length=2048):
+    result = tokenizer(
+        format_chat_history_v2(prompt),
+        truncation=True,
+        max_length=max_length,
+        padding="max_length",
+    )
+    result["labels"] = result["input_ids"]
+    return result
+
+tokenized_dataset = dataset.map(generate_and_tokenize_prompt)
+print(tokenized_dataset['train'][0])
+
+# # stats data length
+# def plot_data_lengths(tokenized_dataset):
+#     lengths = []
+#     for split in tokenized_dataset:
+#         lengths += [len(x['input_ids']) for x in tokenized_dataset[split]]
+#     print(f"Max length: {max(lengths)}")
+#     print(f"Min length: {min(lengths)}")
+#     print(f"Mean length: {sum(lengths)/len(lengths)}")
+#     print(f"Median length: {sorted(lengths)[len(lengths)//2]}")
+
+# plot_data_lengths(tokenized_dataset)
+print(tokenized_dataset['train'][0])
+
+#Adding the adapters in the layers
+from peft import LoraConfig, get_peft_model
+def print_trainable_parameters(model):
+    """
+    Prints the number of trainable parameters in the model.
+    """
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        all_param += param.numel()
+        if param.requires_grad:
+            trainable_params += param.numel()
+    print(
+        f"trainable params: {trainable_params:,} || all params: {all_param:,} || trainable%: {100 * trainable_params / all_param}"
+    )
+model = AutoModelForCausalLM.from_pretrained(base_model_id, torch_dtype=torch.bfloat16)
+# model = prepare_model_for_kbit_training(model)
+peft_config = LoraConfig(
+        r=64,
+        lora_alpha=64,
+        lora_dropout=0.05,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj","gate_proj", "up_proj", "down_proj"]
+    )
+model = get_peft_model(model, peft_config)
+print_trainable_parameters(model)
+print(model)
+
+import wandb, os
+# wandb.login()
+
+wandb_project = "NikaidoLM"
+if len(wandb_project) > 0:
+    os.environ["WANDB_PROJECT"] = wandb_project
+
+import transformers
+from datetime import datetime
+
+project = wandb_project
+base_model_name = "mistral"
+run_name = base_model_name + "-" + project
+output_name = f"{run_name}-{WORK}-{datetime.now().strftime('%Y-%m-%d-%H-%M')}"
+output_dir = "/scratch/generalvision/mowentao/mistral-out/" + output_name
+
+trainer = transformers.Trainer(
+    model=model,
+    train_dataset=tokenized_dataset["train"],
+    eval_dataset=tokenized_dataset["eval"],
+    args=transformers.TrainingArguments(
+        output_dir=output_dir,
+        warmup_steps=500,
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=2,
+        num_train_epochs=3,
+        weight_decay=5e-4,
+        # max_steps=10_000,
+        learning_rate=1e-4, # Want a small lr for finetuning
+        bf16=True,
+        optim="paged_adamw_32bit",
+        logging_steps=100,              # When to start reporting loss
+        logging_dir=output_dir,        # Directory for storing logs
+        save_strategy="steps",       # Save the model checkpoint every logging step
+        save_steps=500,                # Save checkpoints every 50 steps
+        evaluation_strategy="steps", # Evaluate the model every logging step
+        eval_steps=100,               # Evaluate and save checkpoints every 50 steps
+        do_eval=True,                # Perform evaluation at the end of training
+        report_to="wandb",           # Comment this out if you don't want to use weights & baises
+        run_name=output_name,         # Name of the W&B run (optional)
+        lr_scheduler_type="cosine",
+    ),
+    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
+)
+
+model.config.use_cache = False  # silence the warnings. Please re-enable for inference!
+trainer.train()
+trainer.model.save_pretrained(new_model_id)
+wandb.finish()