add train and run scripts

finetune/Llama2_13b/data/eval/tf-ltl_eng_test_mid_ascii_gptAuged/data-00000-of-00001.arrow

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

finetune/Llama2_13b/data/eval/tf-ltl_eng_test_mid_ascii_gptAuged/dataset_info.json

@@ -0,0 +1,69 @@
+{
+  "builder_name": "json",
+  "citation": "",
+  "config_name": "default",
+  "dataset_name": "json",
+  "dataset_size": 889411,
+  "description": "",
+  "download_checksums": {
+    "LTL_datasets/collect/ltl_eng_train_mid_ascii_gptAuged.jsonl": {
+      "num_bytes": 1129386,
+      "checksum": null
+    },
+    "LTL_datasets/collect/ltl_eng_test_mid_ascii_gptAuged.jsonl": {
+      "num_bytes": 125920,
+      "checksum": null
+    }
+  },
+  "download_size": 1255306,
+  "features": {
+    "id": {
+      "dtype": "string",
+      "_type": "Value"
+    },
+    "input_ids": {
+      "feature": {
+        "dtype": "int32",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "attention_mask": {
+      "feature": {
+        "dtype": "int8",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "labels": {
+      "feature": {
+        "dtype": "int64",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    }
+  },
+  "homepage": "",
+  "license": "",
+  "size_in_bytes": 2144717,
+  "splits": {
+    "train": {
+      "name": "train",
+      "num_bytes": 800102,
+      "num_examples": 10621,
+      "dataset_name": "json"
+    },
+    "test": {
+      "name": "test",
+      "num_bytes": 89309,
+      "num_examples": 1181,
+      "dataset_name": "json"
+    }
+  },
+  "version": {
+    "version_str": "0.0.0",
+    "major": 0,
+    "minor": 0,
+    "patch": 0
+  }
+}

finetune/Llama2_13b/data/eval/tf-ltl_eng_test_mid_ascii_gptAuged/state.json

@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "c6bf809a7a8f99a6",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": "test"
+}

finetune/Llama2_13b/data/train/tf-ltl_eng_test_mid_ascii_gptAuged/data-00000-of-00001.arrow

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:56f425bd727940e8905c8c73e87d9089a94304a1627835291e5595f54d984e0f
+size 11945016

finetune/Llama2_13b/data/train/tf-ltl_eng_test_mid_ascii_gptAuged/dataset_info.json

@@ -0,0 +1,69 @@
+{
+  "builder_name": "json",
+  "citation": "",
+  "config_name": "default",
+  "dataset_name": "json",
+  "dataset_size": 889411,
+  "description": "",
+  "download_checksums": {
+    "LTL_datasets/collect/ltl_eng_train_mid_ascii_gptAuged.jsonl": {
+      "num_bytes": 1129386,
+      "checksum": null
+    },
+    "LTL_datasets/collect/ltl_eng_test_mid_ascii_gptAuged.jsonl": {
+      "num_bytes": 125920,
+      "checksum": null
+    }
+  },
+  "download_size": 1255306,
+  "features": {
+    "id": {
+      "dtype": "string",
+      "_type": "Value"
+    },
+    "input_ids": {
+      "feature": {
+        "dtype": "int32",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "attention_mask": {
+      "feature": {
+        "dtype": "int8",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "labels": {
+      "feature": {
+        "dtype": "int64",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    }
+  },
+  "homepage": "",
+  "license": "",
+  "size_in_bytes": 2144717,
+  "splits": {
+    "train": {
+      "name": "train",
+      "num_bytes": 800102,
+      "num_examples": 10621,
+      "dataset_name": "json"
+    },
+    "test": {
+      "name": "test",
+      "num_bytes": 89309,
+      "num_examples": 1181,
+      "dataset_name": "json"
+    }
+  },
+  "version": {
+    "version_str": "0.0.0",
+    "major": 0,
+    "minor": 0,
+    "patch": 0
+  }
+}

finetune/Llama2_13b/data/train/tf-ltl_eng_test_mid_ascii_gptAuged/state.json

@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "afb9c85014ff4b4e",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": "train"
+}

finetune/Llama2_13b/llama_dp2_patch.py

@@ -0,0 +1,139 @@
+from typing import List, Optional, Tuple
+
+import torch
+from torch import nn
+import warnings
+import transformers
+from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
+from peft.tuners.lora import LoraLayer
+
+try:
+    from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func
+    from flash_attn.bert_padding import unpad_input, pad_input
+except Exception:
+    raise ModuleNotFoundError(
+        "Please install FlashAttention first, e.g., with pip install flash-attn --no-build-isolation, Learn more at https://github.com/Dao-AILab/flash-attention#installation-and-features"
+    )
+
+try:
+    from einops import rearrange
+except Exception:
+    raise ModuleNotFoundError("Please install einops first, e.g., with pip install einops")
+
+
+# ADAPTED from https://github.com/allenai/open-instruct/blob/main/open_instruct/llama_flash_attn_monkey_patch.py
+# AND https://github.com/lm-sys/FastChat/blob/main/fastchat/train/llama_flash_attn_monkey_patch.py
+# AND https://github.com/LAION-AI/Open-Assistant/blob/04fa9a24b2a58c8885b8aa6a2eb02b18de6b4961/model/model_training/models/patching_llama.py
+# AND Sourabh https://github.com/huggingface/transformers/commit/ee81bf5aee0d65f005d157c013777e3d27d8d6bf
+def forward(
+    self,
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.Tensor] = None,
+    past_key_value: Optional[Tuple[torch.Tensor]] = None,
+    output_attentions: bool = False,
+    use_cache: bool = False,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    """Input shape: Batch x Time x Channel
+
+    attention_mask: [bsz, q_len]
+    """
+    if output_attentions:
+        warnings.warn("Output attentions is not supported for patched `LlamaAttention`, returning `None` instead.")
+
+    bsz, q_len, _ = hidden_states.size()
+
+    query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    # [bsz, q_len, nh, hd]
+    # [bsz, nh, q_len, hd]
+
+    kv_seq_len = key_states.shape[-2]
+    if past_key_value is not None:
+        kv_seq_len += past_key_value[0].shape[-2]
+    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+    # Past Key value support
+    if past_key_value is not None:
+        # reuse k, v, self_attention
+        key_states = torch.cat([past_key_value[0], key_states], dim=2)
+        value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+    past_key_value = (key_states, value_states) if use_cache else None
+
+    # Flash attention codes from
+    # https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/flash_attention.py
+
+    # transform the data into the format required by flash attention
+    qkv = torch.stack([query_states, key_states, value_states], dim=2)  # [bsz, nh, 3, q_len, hd]
+    qkv = qkv.transpose(1, 3)  # [bsz, q_len, 3, nh, hd]
+    # We have disabled _prepare_decoder_attention_mask in LlamaModel
+    # the attention_mask should be the same as the key_padding_mask
+    key_padding_mask = attention_mask
+
+    if key_padding_mask is None:
+        qkv = rearrange(qkv, "b s ... -> (b s) ...")
+        max_s = q_len
+        cu_q_lens = torch.arange(0, (bsz + 1) * q_len, step=q_len, dtype=torch.int32, device=qkv.device)
+        output = flash_attn_varlen_qkvpacked_func(qkv, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True)
+        output = rearrange(output, "(b s) ... -> b s ...", b=bsz)
+    else:
+        nheads = qkv.shape[-2]
+        x = rearrange(qkv, "b s three h d -> b s (three h d)")
+        x_unpad, indices, cu_q_lens, max_s = unpad_input(x, key_padding_mask)
+        x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads)
+        output_unpad = flash_attn_varlen_qkvpacked_func(
+            x_unpad, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True
+        )
+        output = rearrange(
+            pad_input(rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, bsz, q_len),
+            "b s (h d) -> b s h d",
+            h=nheads,
+        )
+    return self.o_proj(rearrange(output, "b s h d -> b s (h d)")), None, past_key_value
+
+
+# Disable the transformation of the attention mask in LlamaModel as the flash attention
+# requires the attention mask to be the same as the key_padding_mask
+def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
+    # [bsz, seq_len]
+    return attention_mask
+
+
+def replace_attn_with_flash_attn():
+    cuda_major, cuda_minor = torch.cuda.get_device_capability()
+    if cuda_major < 8:
+        print(
+            "Flash attention is only supported on Ampere or Hopper GPU during training due to head dim > 64 backward."
+            "ref: https://github.com/HazyResearch/flash-attention/issues/190#issuecomment-1523359593"
+        )
+    transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = (
+        _prepare_decoder_attention_mask
+    )
+    transformers.models.llama.modeling_llama.LlamaAttention.forward = forward
+
+
+def unplace_flash_attn_with_attn():
+    import importlib
+    import transformers
+
+    print("Reloading llama model, unpatching flash attention")
+    importlib.reload(transformers.models.llama.modeling_llama)
+
+
+# Adapted from https://github.com/tmm1/axolotl/blob/2eda9e02a9d15a7a3f92b41f257d9844d72fc220/src/axolotl/utils/models.py#L338
+def upcast_layer_for_flash_attention(model, torch_dtype):
+    # LlamaRMSNorm layers are in fp32 after kbit_training, so we need to
+    # convert them back to fp16/bf16 for flash-attn compatibility.
+    for name, module in model.named_modules():
+        if isinstance(module, LoraLayer):
+            module.to(torch_dtype)
+        if "norm" in name:
+            module.to(torch_dtype)
+        if "lm_head" in name or "embed_tokens" in name:
+            if hasattr(module, "weight"):
+                module.to(torch_dtype)
+
+    return model

finetune/Llama2_13b/llama_lora_fintune.py

@@ -0,0 +1,264 @@
+from huggingface_hub import login
+login()
+import sys,os
+from datasets import load_dataset
+import torch
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizer, TrainingArguments
+from peft import LoraConfig
+from trl import SFTTrainer
+from accelerate import infer_auto_device_map,init_empty_weights
+
+def preprocess_function(sample,padding="max_length"):
+    # add prefix to the input for t5
+    # print(sample[0])
+    inputs=[
+         f"""### Instruction:
+translate natural description in to LTL, first translate into a logical way, and then translate into LTL,
+using 'A' for 'And','O' for 'Or', 'I' for 'Imply','N' for 'Not','E' for 'Equally','F' for 'Finally','G' for 'Globally','U' for 'Until','X' for 'Next', pay specific attention to brackets '()'
+
+### Input:
+{sample['natural'][i]}
+
+### Response:
+{util.reAsciiLTL2EngLTL(sample['raw_ltl'][i])} {sample['raw_ltl'][i]}</s>"""
+    # NOTE it seems the eos is needed, the bos is not needed(the bos will be automatically added) 
+    for i in (range(len(sample['natural'])))]
+    # inputs = ["## [instruction]: translate natural description in to LTL: ### [natural language]:" + sample['natural'][i]+'### [LTL]:'+sample['raw_ltl'][i] for i in (range(len(sample['natural'])))]
+
+    sample["complete_text"] = inputs
+    return sample
+
+
+def evaluate_model(input_text):
+    input_text =f"""### [Instruction]:
+translate natural description in to LTL, first translate into a logical way, and then translate into LTL,
+using 'A' for 'And','O' for 'Or', 'I' for 'Imply','N' for 'Not','E' for 'Equally','F' for 'Finally','G' for 'Globally','U' for 'Until','X' for 'Next', pay specific attention to brackets '()'
+
+### [Natural Language Task]:
+{input_text}
+### [Temporal Logic Translation]:
+"""
+    # "### [instruction]: translate natural description in to LTL: \n\n ### [natural language]:" + input_text+'### [LTL]:'
+    inputs = tokenizer(input_text, return_tensors="pt").to(device)
+    outputs = model.generate(input_ids=inputs["input_ids"].to("cuda"), attention_mask=inputs["attention_mask"].to("cuda"), max_new_tokens=64, pad_token_id=tokenizer.eos_token_id)
+
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+
+# sys.path.append('../../../')
+# sys.path.append('../../')
+# sys.path.append('../')
+# os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+# os.environ['CUDA_VISIBLE_DEVICES'] = "5,6,7" 
+# device = torch.device("cuda:0-6" if torch.cuda.is_available() else "cpu")
+sys.path.append(os.path.join(os.path.dirname(__file__), '../../'))
+import utils.util as util
+# Load dataset from the hub
+# dataset = load_dataset("samsum")
+output_dir = "finetuned_model/"
+datapath='LTL_datasets/collect/'
+exp_name="ascii"
+
+
+base_model_name = "meta-llama/Llama-2-13b-hf"
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16,
+)
+
+import os
+os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
+os.environ['CUDA_VISIBLE_DEVICES']='4'
+device_map="auto"
+# torch.cuda.set_device(7)
+# device_map={'':torch.cuda.current_device()}
+# device_map = {'':'cuda:7'}
+# model_dir为模型的路径或名称
+# config = AutoConfig.from_pretrained(base_model_name, trust_remote_code=True)
+# with init_empty_weights():
+#     base_model = AutoModelForCausalLM.from_pretrained(
+#     base_model_name,
+#     from_tf=bool(".ckpt" in base_model_name),
+#     quantization_config=bnb_config,
+#     device_map=device_map,
+#     trust_remote_code=True,
+#     use_auth_token=True
+#     )
+
+# map_list = {5:"15GB", 6:"15GB",7:"15GB"}    # 对应不同卡号限制的内存量
+# map_list = {7:"15GB",}    # 对应不同卡号限制的内存量
+# no_split_modules = base_model._no_split_modules
+# device_map = infer_auto_device_map(base_model, max_memory=map_list, no_split_module_classes=no_split_modules)
+
+
+tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
+# tokenizer = AutoTokenizer.from_pretrained(base_model_name, add_eos_token=True,trust_remote_code=True)
+# NOTE no one says whether the add eos token need to be added, but if we do not add this, the generate will continue until reach the max_new_tokens
+# when add add_eos_token, it always failed 
+# if use this it will generate somthing other
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = 'right'
+print(tokenizer.eos_token_id)
+# 2
+print(tokenizer.bos_token_id)
+# 1
+# print(tokenizer._convert_token_to_id(tokenizer.bos_token))
+
+dataset = load_dataset("json",  data_files={"train":"LTL_datasets/collect/ltl_eng_train_mid_ascii_gptAuged.jsonl","test":"LTL_datasets/collect/ltl_eng_test_mid_ascii_gptAuged.jsonl"})
+print(dataset)
+
+from datasets import concatenate_datasets
+import numpy as np
+# The maximum total input sequence length after tokenization.
+# Sequences longer than this will be truncated, sequences shorter will be padded.
+# tokenized_inputs = concatenate_datasets([dataset["train"], dataset["test"]]).map(lambda x: tokenizer(x["natural"], truncation=True), batched=True, remove_columns=["raw_ltl", "natural"])
+# input_lenghts = [len(x) for x in tokenized_inputs["input_ids"]]
+# # take 85 percentile of max length for better utilization
+# max_source_length = int(np.percentile(input_lenghts, 100))
+# print(f"Max source length: {max_source_length}")
+
+# # The maximum total sequence length for target text after tokenization.
+# # Sequences longer than this will be truncated, sequences shorter will be padded."
+# tokenized_targets = concatenate_datasets([dataset["train"], dataset["test"]]).map(lambda x: tokenizer(x["raw_ltl"], truncation=True), batched=True, remove_columns=["raw_ltl", "natural"])
+# target_lenghts = [len(x) for x in tokenized_targets["input_ids"]]
+# # take 90 percentile of max length for better utilization
+# max_target_length = int(np.percentile(target_lenghts, 100))
+# print(f"Max target length: {max_target_length}")
+
+
+# # %%
+# def translateAscii2Eng(input):
+
+# def preprocess_function(sample,padding="max_length"):
+#     # add prefix to the input for t5
+#     # print(sample[0])
+#     inputs = ["## [instruction]: translate natural description in to LTL: ### [natural language]:" + sample['natural'][i]+'### [LTL]:'+sample['raw_ltl'][i] for i in (range(len(sample['natural'])))]
+
+#     sample["complete_text"] = inputs
+#     return sample
+
+
+
+tokenized_dataset = dataset.map(preprocess_function, batched=True)
+print(f"Keys of tokenized dataset: {list(tokenized_dataset['train'].features)}")
+
+# save datasets to disk for later easy loading
+# tokenized_dataset["train"].save_to_disk("data/train"+exp_name)
+# tokenized_dataset["test"].save_to_disk("data/eval"+exp_name)
+
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
+    from_tf=bool(".ckpt" in base_model_name),
+    quantization_config=bnb_config,
+    device_map=device_map,
+    trust_remote_code=True,
+    use_auth_token=True
+)
+base_model.config.use_cache = False
+
+# More info: https://github.com/huggingface/transformers/pull/24906
+base_model.config.pretraining_tp = 1 
+
+peft_config = LoraConfig(
+    lora_alpha=16,
+    lora_dropout=0.1,
+    r=64,
+    bias="none",
+    task_type="CAUSAL_LM",
+)
+
+
+
+
+training_args = TrainingArguments(
+    output_dir=output_dir,
+    per_device_train_batch_size=8,
+    gradient_accumulation_steps=4,
+    learning_rate=2e-4,
+    logging_steps=10,
+    num_train_epochs=3,
+    # max_steps=500
+)
+
+max_seq_length = 512
+
+trainer = SFTTrainer(
+    model=base_model,
+    train_dataset=tokenized_dataset['train'],
+    peft_config=peft_config,
+    dataset_text_field="complete_text",
+    max_seq_length=max_seq_length,
+    tokenizer=tokenizer,
+    args=training_args,
+    # device_map=device_map
+)
+
+import os
+output_dir = os.path.join(output_dir, "llama2_13b_"+exp_name+'aug1')
+
+trainer.train()
+trainer.model.save_pretrained(output_dir)
+# trainer.model.save_pretrained(output_dir)
+tokenizer.save_pretrained(output_dir)
+
+
+# check
+
+from peft import AutoPeftModelForCausalLM
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map=device_map, torch_dtype=torch.bfloat16)
+
+
+
+
+import evaluate
+import numpy as np
+from datasets import load_from_disk
+from tqdm import tqdm
+
+# Metric
+metric = evaluate.load("rouge")
+
+
+# load test dataset from distk
+# test_dataset = load_from_disk("data/eval"+exp_name+'/').with_format("torch")
+
+
+# run predictions
+# this can take ~45 minutes
+import re 
+pattern=re.compile("\[Formal LTL\]:\n([\S ]*)\n")
+predictions, references,input_sentence,output_sentence=[], [] , [], []
+for idx in range(len(tokenized_dataset['test']['natural'])):
+    # print(sample)
+    nl=tokenized_dataset['test']['natural'][idx]
+    p = evaluate_model(nl)
+    # print(p,l)
+    input_sentence.append(nl)
+
+    transLTL=pattern.findall(p)
+    # print(p)
+    predictions.append(transLTL[0]) 
+    output_sentence.append(p) 
+    input_sentence.append(p)
+    references.append(tokenized_dataset['test']['raw_ltl'][idx])
+    print(input_sentence[-1],'\nout::\n',output_sentence[-1],'\npre::\n',predictions[-1],'\nref::\n',references[-1],'\n','-'*20,'\n')
+
+# compute metric
+rogue = metric.compute(predictions=predictions, references=references, use_stemmer=True)
+
+# print results
+print(f"Rogue1: {rogue['rouge1']* 100:2f}%")
+print(f"rouge2: {rogue['rouge2']* 100:2f}%")
+print(f"rougeL: {rogue['rougeL']* 100:2f}%")
+print(f"rougeLsum: {rogue['rougeLsum']* 100:2f}%")
+eval_output=np.array([input_sentence,predictions,references]).T
+import pandas as pd 
+eval_output=pd.DataFrame(eval_output)
+pd.DataFrame.to_csv(eval_output,output_dir+'/output')
+# Rogue1: 50.386161%
+# rouge2: 24.842412%
+# rougeL: 41.370130%
+# rougeLsum: 41.394230%

finetune/Llama2_13b/llama_lora_fintune_ver2.py

@@ -0,0 +1,281 @@
+from huggingface_hub import login
+login()
+import json
+import numpy as npas np
+import sys,os
+from datasets import load_dataset
+import torch
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizer, TrainingArguments
+from peft import LoraConfig
+from trl import SFTTrainer
+from accelerate import infer_auto_device_map,init_empty_weights
+# sys.path.append('../../../')
+# sys.path.append('../../')
+# sys.path.append('../')
+# os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+# os.environ['CUDA_VISIBLE_DEVICES'] = "5,6,7" 
+# device = torch.device("cuda:0-6" if torch.cuda.is_available() else "cpu")
+sys.path.append(os.path.join(os.path.dirname(__file__), '../../'))
+# import utils.util as util
+# Load dataset from the hub
+# dataset = load_dataset("samsum")
+np.random.seed(42)
+output_dir = "model_weight/"
+datapath='NL2TL-dataset/collect2'
+exp_name="ascii"
+explainer_files=['LTLexplain_0.json','LTLexplain_1.json','LTLexplain_2.json','LTLexplain_3.json']
+explainer_dic={}
+for path in explainer_files:
+    with open(os.path.join(datapath,path)) as f:
+        LTLlist=json.load(f)
+    for key in LTLlist.keys():
+        if isinstance(LTLlist[key],dict):
+            if not (key in explainer_dic):
+                explainer_dic[key]=[]
+            explainer_dic[key].append(LTLlist[key]['translate'])
+            sp=LTLlist[key]['explain'].split("means that")
+            if len(sp)>1:
+                explainer_dic[key].append(sp[1])
+    
+
+
+base_model_name = "meta-llama/Llama-2-13b-hf"
+base_model_name = "meta-llama/Llama-2-7b-hf"
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16,
+)
+
+import os
+os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
+os.environ['CUDA_VISIBLE_DEVICES']='4'
+device_map="auto"
+# torch.cuda.set_device(7)
+device_map={'':torch.cuda.current_device()}
+device_map={'':torch.cuda.current_device()}
+# device_map = {'':'cuda:7'}
+# model_dir为模型的路径或名称
+# config = AutoConfig.from_pretrained(base_model_name, trust_remote_code=True)
+# with init_empty_weights():
+#     base_model = AutoModelForCausalLM.from_pretrained(
+#     base_model_name,
+#     from_tf=bool(".ckpt" in base_model_name),
+#     quantization_config=bnb_config,
+#     device_map=device_map,
+#     trust_remote_code=True,
+#     use_auth_token=True
+#     )
+
+# map_list = {5:"15GB", 6:"15GB",7:"15GB"}    # 对应不同卡号限制的内存量
+# map_list = {7:"15GB",}    # 对应不同卡号限制的内存量
+# no_split_modules = base_model._no_split_modules
+# device_map = infer_auto_device_map(base_model, max_memory=map_list, no_split_module_classes=no_split_modules)
+
+
+tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
+# tokenizer = AutoTokenizer.from_pretrained(base_model_name, add_eos_token=True,trust_remote_code=True)
+# NOTE no one says whether the add eos token need to be added, but if we do not add this, the generate will continue until reach the max_new_tokens
+# when add add_eos_token, it always failed 
+# if use this it will generate somthing other
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = 'right'
+print(tokenizer.eos_token_id)
+# 2
+print(tokenizer.bos_token_id)
+# 1
+# print(tokenizer._convert_token_to_id(tokenizer.bos_token))
+
+dataset = load_dataset("json",  data_files={"train":os.path.join(datapath,"ltl_eng_train_mid_ascii_gptAuged.jsonl"),"test":os.path.join(datapath,"ltl_eng_test_mid_ascii_gptAuged.jsonl")})
+print(dataset)
+
+from datasets import concatenate_datasets
+import numpy as np
+# The maximum total input sequence length after tokenization.
+# Sequences longer than this will be truncated, sequences shorter will be padded.
+# tokenized_inputs = concatenate_datasets([dataset["train"], dataset["test"]]).map(lambda x: tokenizer(x["natural"], truncation=True), batched=True, remove_columns=["raw_ltl", "natural"])
+# input_lenghts = [len(x) for x in tokenized_inputs["input_ids"]]
+# # take 85 percentile of max length for better utilization
+# max_source_length = int(np.percentile(input_lenghts, 100))
+# print(f"Max source length: {max_source_length}")
+
+# # The maximum total sequence length for target text after tokenization.
+# # Sequences longer than this will be truncated, sequences shorter will be padded."
+# tokenized_targets = concatenate_datasets([dataset["train"], dataset["test"]]).map(lambda x: tokenizer(x["raw_ltl"], truncation=True), batched=True, remove_columns=["raw_ltl", "natural"])
+# target_lenghts = [len(x) for x in tokenized_targets["input_ids"]]
+# # take 90 percentile of max length for better utilization
+# max_target_length = int(np.percentile(target_lenghts, 100))
+# print(f"Max target length: {max_target_length}")
+
+
+# # %%
+# def translateAscii2Eng(input):
+
+# def preprocess_function(sample,padding="max_length"):
+#     # add prefix to the input for t5
+#     # print(sample[0])
+#     inputs = ["## [instruction]: translate natural description in to LTL: ### [natural language]:" + sample['natural'][i]+'### [LTL]:'+sample['raw_ltl'][i] for i in (range(len(sample['natural'])))]
+
+#     sample["complete_text"] = inputs
+#     return sample
+
+
+def preprocess_function(sample,padding="max_length"):
+    # add prefix to the input for t5
+    # print(sample[0])
+    inputs=[
+         f"""### Instruction:
+translate natural description to linear temproal logic, first translate into a logical way, and then translate into linear temproal logic, pay specific attention to brackets '()'
+
+### Natural Language Task:
+{sample['natural'][i].strip()}
+
+### Logic Translation:
+{explainer_dic[sample['raw_ltl'][i].strip()][np.random.randint(0,len(explainer_dic[sample['raw_ltl'][i].strip()]))]}
+
+### linear temproal logic:
+{sample['raw_ltl'][i].strip()}
+</s>""".lower()
+    # NOTE it seems the eos is needed, the bos is not needed(the bos will be automatically added) 
+    for i in (range(len(sample['natural'])))]
+    # inputs = ["## [instruction]: translate natural description in to LTL: ### [natural language]:" + sample['natural'][i]+'### [LTL]:'+sample['raw_ltl'][i] for i in (range(len(sample['natural'])))]
+
+    sample["complete_text"] = inputs
+    return sample
+
+tokenized_dataset = dataset.map(preprocess_function, batched=True)
+print(f"Keys of tokenized dataset: {list(tokenized_dataset['train'].features)}")
+
+# save datasets to disk for later easy loading
+# tokenized_dataset["train"].save_to_disk("data/train"+exp_name)
+# tokenized_dataset["test"].save_to_disk("data/eval"+exp_name)
+
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
+    from_tf=bool(".ckpt" in base_model_name),
+    # quantization_config=bnb_config,
+    device_map=device_map,
+    trust_remote_code=True,
+    use_auth_token=True
+)
+base_model.config.use_cache = False
+
+# More info: https://github.com/huggingface/transformers/pull/24906
+base_model.config.pretraining_tp = 1 
+
+peft_config = LoraConfig(
+    lora_alpha=16,
+    lora_dropout=0.1,
+    r=64,
+    bias="none",
+    task_type="CAUSAL_LM",
+)
+
+
+
+
+training_args = TrainingArguments(
+    output_dir=output_dir,
+    per_device_train_batch_size=8,
+    gradient_accumulation_steps=4,
+    learning_rate=2e-4,
+    logging_steps=10,
+    num_train_epochs=3,
+    # max_steps=500
+)
+
+max_seq_length = 512
+
+trainer = SFTTrainer(
+    model=base_model,
+    train_dataset=tokenized_dataset['train'],
+    peft_config=peft_config,
+    dataset_text_field="complete_text",
+    max_seq_length=max_seq_length,
+    tokenizer=tokenizer,
+    args=training_args,
+    # device_map=device_map
+)
+
+import os
+output_dir = os.path.join(output_dir, "llama2_13b_"+exp_name+'aug1')
+
+trainer.train()
+trainer.model.save_pretrained(output_dir)
+# trainer.model.save_pretrained(output_dir)
+tokenizer.save_pretrained(output_dir)
+
+
+# check
+
+from peft import AutoPeftModelForCausalLM
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map=device_map, torch_dtype=torch.bfloat16)
+
+
+def evaluate_model(input_text):
+    input_text =f"""### ### Instruction:
+translate natural description to linear temproal logic, first translate into a logical way, and then translate into linear temproal logic, pay specific attention to brackets '()'
+
+### Natural Language Task:
+{input_text}
+
+"""
+    # "### [instruction]: translate natural description in to LTL: \n\n ### [natural language]:" + input_text+'### [LTL]:'
+    inputs = tokenizer(input_text, return_tensors="pt").to(device)
+    outputs = model.generate(input_ids=inputs["input_ids"].to("cuda"), attention_mask=inputs["attention_mask"].to("cuda"), max_new_tokens=64, pad_token_id=tokenizer.eos_token_id)
+
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# if __name__=='__main__':
+import evaluate
+import numpy as np
+from datasets import load_from_disk
+from tqdm import tqdm
+
+# Metric
+metric = evaluate.load("rouge")
+
+
+# load test dataset from distk
+# test_dataset = load_from_disk("data/eval"+exp_name+'/').with_format("torch")
+
+
+# run predictions
+# this can take ~45 minutes
+import re 
+pattern=re.compile("### linear temproal logic::\n([\S ]*)\n")
+predictions, references,input_sentence,output_sentence=[], [] , [], []
+for idx in range(len(tokenized_dataset['test']['natural'])):
+    # print(sample)
+    nl=tokenized_dataset['test']['natural'][idx]
+    p = evaluate_model(nl)
+    # print(p,l)
+    input_sentence.append(nl)
+
+    transLTL=pattern.findall(p)
+    # print(p)
+    predictions.append(transLTL[0].strip()) 
+    output_sentence.append(p) 
+    input_sentence.append(p)
+    references.append(tokenized_dataset['test']['raw_ltl'][idx].strip())
+    references.append(tokenized_dataset['test']['raw_ltl'][idx].strip())
+    print(input_sentence[-1],'\nout::\n',output_sentence[-1],'\npre::\n',predictions[-1],'\nref::\n',references[-1],'\n','-'*20,'\n')
+
+# compute metric
+rogue = metric.compute(predictions=predictions, references=references, use_stemmer=True)
+
+# print results
+print(f"Rogue1: {rogue['rouge1']* 100:2f}%")
+print(f"rouge2: {rogue['rouge2']* 100:2f}%")
+print(f"rougeL: {rogue['rougeL']* 100:2f}%")
+print(f"rougeLsum: {rogue['rougeLsum']* 100:2f}%")
+eval_output=np.array([input_sentence,predictions,references]).T
+import pandas as pd 
+eval_output=pd.DataFrame(eval_output)
+pd.DataFrame.to_csv(eval_output,output_dir+'/output')
+# Rogue1: 50.386161%
+# rouge2: 24.842412%
+# rougeL: 41.370130%
+# rougeLsum: 41.394230%

finetune/Llama2_13b/llama_lora_fintune_ver3_qlora.py

@@ -0,0 +1,336 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+# device = "cuda" # the device to load the model onto
+# from huggingface_hub import login
+# login()
+import json
+import numpy as np
+import sys,os
+from datasets import load_dataset
+import torch
+from transformers import (AutoModelForCausalLM,
+                         AutoTokenizer,
+                         BitsAndBytesConfig,
+                         TrainingArguments,
+                         pipeline,
+                         logging,
+                         TrainerCallback)
+from peft import LoraConfig, PeftConfig, prepare_model_for_kbit_training, get_peft_model
+from trl import SFTTrainer
+from accelerate import infer_auto_device_map,init_empty_weights
+import wandb
+from datasets import concatenate_datasets
+import numpy as np
+# sys.path.append('../../../')
+# sys.path.append('../../')
+# sys.path.append('../')
+# os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+# os.environ['CUDA_VISIBLE_DEVICES'] = "5,6,7" 
+# device = torch.device("cuda:0-6" if torch.cuda.is_available() else "cpu")
+sys.path.append(os.path.join(os.path.dirname(__file__), '../../'))
+# import utils.util as util
+# Load dataset from the hub
+# dataset = load_dataset("samsum")
+np.random.seed(42)
+output_dir = "model_weight/"
+datapath='NL2TL-dataset/collect2'
+exp_name="ascii"
+explainer_files=['LTLexplain_0.json','LTLexplain_1.json','LTLexplain_2.json','LTLexplain_3.json']
+explainer_dic={}
+for path in explainer_files:
+    with open(os.path.join(datapath,path)) as f:
+        LTLlist=json.load(f)
+    for key in LTLlist.keys():
+        if isinstance(LTLlist[key],dict):
+            if not (key in explainer_dic):
+                explainer_dic[key]=[]
+            explainer_dic[key].append(LTLlist[key]['translate'])
+            sp=LTLlist[key]['explain'].split("means that")
+            if len(sp)>1:
+                explainer_dic[key].append(sp[1])
+    
+base_model_name = "meta-llama/Llama-2-13b-hf"
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit = True,
+    bnb_4bit_use_double_quant = False,
+    bnb_4bit_quant_type = 'nf4',
+    bnb_4bit_compute_dtype = getattr(torch, "float16")
+)
+
+import os
+os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
+os.environ['CUDA_VISIBLE_DEVICES']='0'
+device_map="auto"
+# torch.cuda.set_device(7)
+# device_map={'':torch.cuda.current_device()}
+# device_map = {'':'cuda:7'}
+# model_dir为模型的路径或名称
+# config = AutoConfig.from_pretrained(base_model_name, trust_remote_code=True)
+# with init_empty_weights():
+#     base_model = AutoModelForCausalLM.from_pretrained(
+#     base_model_name,
+#     from_tf=bool(".ckpt" in base_model_name),
+#     quantization_config=bnb_config,
+#     device_map=device_map,
+#     trust_remote_code=True,
+#     use_auth_token=True
+#     )
+
+# map_list = {5:"15GB", 6:"15GB",7:"15GB"}    # 对应不同卡号限制的内存量
+# map_list = {7:"15GB",}    # 对应不同卡号限制的内存量
+# no_split_modules = base_model._no_split_modules
+# device_map = infer_auto_device_map(base_model, max_memory=map_list, no_split_module_classes=no_split_modules)
+
+
+dataset = load_dataset("json",  data_files={"train":os.path.join(datapath,"ltl_eng_train_mid_ascii_gptAuged.jsonl"),"test":os.path.join(datapath,"ltl_eng_test_mid_ascii_gptAuged.jsonl")})
+print(dataset)
+
+
+
+tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
+# tokenizer = AutoTokenizer.from_pretrained(base_model_name, add_eos_token=True,trust_remote_code=True)
+# NOTE no one says whether the add eos token need to be added, but if we do not add this, the generate will continue until reach the max_new_tokens
+# when add add_eos_token, it always failed 
+# if use this it will generate somthing other
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = 'right'
+print(tokenizer.eos_token_id)
+# 2
+print(tokenizer.bos_token_id)
+# 1
+# print(tokenizer._convert_token_to_id(tokenizer.bos_token))
+
+def preprocess_function(sample,padding="max_length"):
+    # add prefix to the input for t5
+    # print(sample[0])
+    inputs=[
+         f"""### Instruction:
+translate natural description to linear temproal logic, first translate into a logical way, and then translate into linear temproal logic, pay specific attention to brackets '()'
+
+### Natural Language Task:
+{sample['natural'][i].strip()}
+
+### Logic Translation:
+{explainer_dic[sample['raw_ltl'][i].strip()][np.random.randint(0,len(explainer_dic[sample['raw_ltl'][i].strip()]))]}
+
+### linear temproal logic:
+{sample['raw_ltl'][i].strip()}
+</s>""".lower()
+    # NOTE it seems the eos is needed, the bos is not needed(the bos will be automatically added) 
+    for i in (range(len(sample['natural'])))]
+    # inputs = ["## [instruction]: translate natural description in to LTL: ### [natural language]:" + sample['natural'][i]+'### [LTL]:'+sample['raw_ltl'][i] for i in (range(len(sample['natural'])))]
+
+    sample["complete_text"] = inputs
+    return sample
+# method1
+# tokenized_dataset = dataset.map(preprocess_function, batched=True)
+# method2
+def preprocess_function2(sample,padding="max_length"):
+    # add prefix to the input for t5
+    # print(sample[0])
+    inputs=[
+         tokenizer.apply_chat_template(
+    [
+        {"role": "system", "content": "translate natural description to linear temproal logic, first translate into a logical way, and then translate into linear temproal logic, pay specific attention to brackets '()'"},
+        {"role": "user", "content": "Natural Language Task: {}".format(sample['natural'][i].strip())},
+        {"role": "assistant", "content": "Logic Translation is {}, linear temproal logic is {}".format(
+            explainer_dic[sample['raw_ltl'][i].strip()][np.random.randint(0,len(explainer_dic[sample['raw_ltl'][i].strip()]))],
+            sample['raw_ltl'][i].strip()
+            )
+        }
+    ],tokenize=False, add_generation_prompt=False)
+    # NOTE it seems the eos is needed, the bos is not needed(the bos will be automatically added) 
+    for i in (range(len(sample['natural'])))]
+    # inputs = ["## [instruction]: translate natural description in to LTL: ### [natural language]:" + sample['natural'][i]+'### [LTL]:'+sample['raw_ltl'][i] for i in (range(len(sample['natural'])))]
+
+    sample["complete_text"] = inputs
+    return sample
+tokenized_dataset = dataset.map(preprocess_function2, batched=True)
+print(f"Keys of tokenized dataset: {list(tokenized_dataset['train'].features)}")
+
+# save datasets to disk for later easy loading
+# tokenized_dataset["train"].save_to_disk("data/train"+exp_name)
+# tokenized_dataset["test"].save_to_disk("data/eval"+exp_name)
+
+class PeftSavingCallback(TrainerCallback):
+    def on_save(self, args, state, control, **kwargs):
+        checkpoint_path = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")
+        kwargs["model"].save_pretrained(checkpoint_path)
+
+        if "pytorch_model.bin" in os.listdir(checkpoint_path):
+            os.remove(os.path.join(checkpoint_path, "pytorch_model.bin"))
+callbacks = [PeftSavingCallback]
+
+peft_config = LoraConfig(
+    lora_alpha=16,
+    lora_dropout=0.05,
+    r=128,
+    bias="none",
+    task_type="CAUSAL_LM",
+    target_modules=["q_proj", "v_proj"]
+)
+
+
+training_arguments = TrainingArguments(
+    output_dir=output_dir,
+    logging_dir = os.path.join(output_dir,"logs"),
+    per_device_train_batch_size=32,
+    num_train_epochs=3,
+    gradient_accumulation_steps=1,
+    optim="paged_adamw_32bit",
+    save_strategy='epoch',
+    logging_steps=25,
+    learning_rate=2e-4,
+    weight_decay=0.001,
+    fp16=True,
+    bf16=False,
+    max_grad_norm=0.3,
+    max_steps=-1,
+    warmup_ratio = 0.05,
+    group_by_length=True,
+    lr_scheduler_type="cosine",
+    report_to="wandb",
+    evaluation_strategy="epoch",
+    do_eval=True,
+    run_name = base_model_name+exp_name,
+    disable_tqdm=False
+)
+import os
+output_dir = os.path.join(output_dir, "llama2_13b"+exp_name+'aug1')
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
+    from_tf=bool(".ckpt" in base_model_name),
+    quantization_config=bnb_config,
+    device_map=device_map,
+    trust_remote_code=True,
+    use_auth_token=True
+)
+base_model.config.use_cache = False
+
+# More info: https://github.com/huggingface/transformers/pull/24906
+base_model.config.pretraining_tp = 1 
+
+base_model.gradient_checkpointing_enable()
+base_model = prepare_model_for_kbit_training(base_model)
+base_model = get_peft_model(base_model, peft_config)
+
+trainer = SFTTrainer(
+    model=base_model,
+    train_dataset=tokenized_dataset['train'],
+    eval_dataset=tokenized_dataset['test'],
+    peft_config=peft_config,
+    dataset_text_field="complete_text",
+    max_seq_length=512,
+    tokenizer=tokenizer,
+    args=training_arguments,
+    callbacks=callbacks,
+    packing=False,
+)
+# wandb.login()
+# trainer.train()
+# trainer.model.save_pretrained(output_dir)
+# # trainer.model.save_pretrained(output_dir)
+# tokenizer.save_pretrained(output_dir)
+
+# wandb.finish()
+
+# check
+print('model dir',output_dir)
+from peft import AutoPeftModelForCausalLM
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = AutoPeftModelForCausalLM.from_pretrained(output_dir,
+    from_tf=bool(".ckpt" in base_model_name),
+    quantization_config=bnb_config,
+    device_map=device_map,
+    trust_remote_code=True,
+    use_auth_token=True
+)
+tokenizer = AutoTokenizer.from_pretrained(output_dir, trust_remote_code=True)
+print(tokenizer.default_chat_template)
+def evaluate_model(input_text):
+    input_text =f"""### Instruction:
+translate natural description to linear temproal logic, first translate into a logical way, and then translate into linear temproal logic, pay specific attention to brackets '()' ### Natural Language Task:
+{input_text}""".lower()
+    inputs = tokenizer(input_text, return_tensors="pt").to(device)
+    print(inputs)
+    outputs = model.generate(input_ids=inputs["input_ids"].to("cuda"), attention_mask=inputs["attention_mask"].to("cuda"), max_new_tokens=512, pad_token_id=tokenizer.eos_token_id)
+
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+def evaluate_model2(input_text):
+    messages=[
+        {"role": "system", "content": "translate natural description to linear temproal logic, first translate into a logical way, and then translate into linear temproal logic, pay specific attention to brackets '()'"},
+        {"role": "user", "content": "Natural Language Task: {}".format(input_text)},
+    ]
+
+    encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt")
+    outputs = model.generate(encodeds, max_new_tokens=512, pad_token_id=tokenizer.eos_token_id)
+#     input_text =f"""### Instruction:
+# translate natural description to linear temproal logic, first translate into a logical way, and then translate into linear temproal logic, pay specific attention to brackets '()' ### Natural Language Task:
+# {input_text}""".lower()
+#     inputs = tokenizer(input_text, return_tensors="pt").to(device)
+#     print(inputs)
+#     outputs = model.generate(input_ids=inputs["input_ids"].to("cuda"), attention_mask=inputs["attention_mask"].to("cuda"), max_new_tokens=512, pad_token_id=tokenizer.eos_token_id)
+
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+# if __name__=='__main__':
+import evaluate
+import numpy as np
+from datasets import load_from_disk
+from tqdm import tqdm
+
+# Metric
+metric = evaluate.load("rouge")
+
+
+# load test dataset from distk
+# test_dataset = load_from_disk("data/eval"+exp_name+'/').with_format("torch")
+
+# run predictions
+# this can take ~45 minutes
+import re 
+pattern=re.compile("linear temproal logic:\n([\S ]*)\n")
+predictions, references,input_sentence,output_sentence=[], [] , [], []
+for idx in range(len(tokenized_dataset['test']['natural'])):
+    # print(sample)
+    nl=tokenized_dataset['test']['natural'][idx]
+    p = evaluate_model2(nl)
+    # print(p,l)
+    input_sentence.append(nl)
+
+    transLTL=pattern.findall(p)
+    print(p)
+    predictions.append(transLTL[0].strip()) 
+    output_sentence.append(p) 
+    input_sentence.append(p)
+    references.append(tokenized_dataset['test']['raw_ltl'][idx].strip())
+    print(input_sentence[-1],'\nout::\n',output_sentence[-1],'\npre::\n',predictions[-1],'\nref::\n',references[-1],'\n','-'*20,'\n')
+
+# compute metric
+rogue = metric.compute(predictions=predictions, references=references, use_stemmer=True)
+
+# print results
+print(f"Rogue1: {rogue['rouge1']* 100:2f}%")
+print(f"rouge2: {rogue['rouge2']* 100:2f}%")
+print(f"rougeL: {rogue['rougeL']* 100:2f}%")
+print(f"rougeLsum: {rogue['rougeLsum']* 100:2f}%")
+eval_output=np.array([input_sentence,predictions,references]).T
+import pandas as pd 
+eval_output=pd.DataFrame(eval_output)
+pd.DataFrame.to_csv(eval_output,output_dir+'/output')
+
+exit()
+messages = [
+    {"role": "user", "content": "What is your favourite condiment?"},
+    {"role": "assistant", "content": "Well, I'm quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever I'm cooking up in the kitchen!"},
+    {"role": "user", "content": "Do you have mayonnaise recipes?"}
+]
+
+encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt")
+
+model_inputs = encodeds.to(device)
+model.to(device)
+
+generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
+decoded = tokenizer.batch_decode(generated_ids)
+print(decoded[0])

finetune/Llama2_13b/llama_lora_test.py

@@ -0,0 +1,185 @@
+# from huggingface_hub import login
+# login()
+import sys,os
+from datasets import load_dataset
+import torch
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizer, TrainingArguments
+# from peft import LoraConfig
+# from trl import SFTTrainer
+# from accelerate import infer_auto_device_map,init_empty_weights
+
+# sys.path.append(os.path.join(os.path.dirname(__file__), '../../'))
+from NL2HLTLtaskPlanner.utils import Task2Preplacer
+from NL2HLTLtaskPlanner.utils import LTLChecker
+import re 
+from datasets import concatenate_datasets
+import numpy as np
+from peft import AutoPeftModelForCausalLM
+os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
+# os.environ['CUDA_VISIBLE_DEVICES']='3'
+
+
+
+class Llama_NL2TL_translator():
+    def __init__(self,
+                 output_dir = "/home/icl-mill19/xsj/model_weight",  
+                 tuned_model_name="llama2_13b__mid_asciiaug1",
+                #  CUDA_device='0',
+                 quat=True) -> None:
+        # os.environ['CUDA_VISIBLE_DEVICES']=CUDA_device
+        self.device_map="auto"
+        self.base_model_name = "meta-llama/Llama-2-13b-hf"
+        self.output_dir = os.path.join(output_dir, tuned_model_name)
+        # check
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        # AutoPeftModelForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf")
+
+
+        # quantconfig = BitsAndBytesConfig(
+        #     load_in_8bit=True,
+        #     bnb_8bit_quant_type="nf4",
+        #     bnb_8bit_use_double_quant=True,
+        #     bnb_8bit_compute_dtype=torch.bfloat16,
+        # )
+        if quat==False:
+            self.model = AutoPeftModelForCausalLM.from_pretrained(self.output_dir, device_map=self.device_map, torch_dtype=torch.bfloat16)
+            # ICL super man可以不量化
+        else:
+            self.model = AutoPeftModelForCausalLM.from_pretrained(self.output_dir,device_map=self.device_map,  torch_dtype=torch.float16,
+            load_in_8bit=True)
+            # quantization_config=quantconfig)
+
+        self.tokenizer = AutoTokenizer.from_pretrained(self.output_dir, trust_remote_code=True)
+        # tokenizer = AutoTokenizer.from_pretrained(base_model_name, add_eos_token=True,trust_remote_code=True)
+        # NOTE no one says whether the add eos token need to be added, but if we do not add this, the generate will continue until reach the max_new_tokens, 
+        # when in predict model, do not use the add_eos_token=True, as the tokenizer will automatically add <\s> to the input, and thus the output will be inregular
+        # when add add_eos_token, it always failed 
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.padding_side = 'right'
+        print(self.tokenizer.eos_token_id)
+        # 2
+        print(self.tokenizer.bos_token_id)
+        # 1
+        # print(tokenizer._convert_token_to_id(tokenizer.bos_token))
+
+        print("NL2TL model loaded")
+        self.pattern=re.compile("\[Formal LTL\]:\n([\S ]*)\n")
+        self.replace=Task2Preplacer()
+        self.ltlChecker=LTLChecker()
+        pass
+
+        print('NL2TL llama translate test:')
+        self.translate("Task_1.1 must be done, and Task_1.2 should be finished before Task_1.1")
+    def evaluate_model(self,input_text):
+        input_text =f"""### [Instruction]:
+translate natural description in to LTL, first translate into a logical way, and then translate into LTL,
+using 'A' for 'And','O' for 'Or', 'I' for 'Imply','N' for 'Not','E' for 'Equally','F' for 'Finally','G' for 'Globally','U' for 'Until','X' for 'Next', pay specific attention to brackets '()'
+
+### [Natural Language Task]:
+{input_text}
+### [Temporal Logic Translation]:
+"""
+        # "### [instruction]: translate natural description in to LTL: \n\n ### [natural language]:" + input_text+'### [LTL]:'
+        inputs = self.tokenizer(input_text, return_tensors="pt").to(self.device)
+        outputs = self.model.generate(input_ids=inputs["input_ids"].to("cuda"), attention_mask=inputs["attention_mask"].to("cuda"), max_new_tokens=128, pad_token_id=self.tokenizer.eos_token_id)
+
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    def translate(self,input_prompt:str=""):
+        print(input_prompt)
+        input_prompt=self.replace.reTask2P(input_prompt)
+        # print(predicter( replace.reTask2P(input_prompt)))
+        # print(input_prompt)
+
+        p=self.evaluate_model(input_prompt)
+        # print(p)
+        transLTL=self.pattern.findall(p)[0]
+        print(transLTL)
+        while(not (self.ltlChecker.AP_CorrCheck(input_prompt,transLTL) and self.ltlChecker.brackets_Check(transLTL))):
+            p=self.evaluate_model(input_prompt)
+            # print(p)
+            transLTL=self.pattern.findall(p)[0]
+            print(transLTL)
+
+        return self.replace.reP2Task(transLTL)
+
+
+if __name__=="__main__":
+    translater=Llama_NL2TL_translator()
+    # test_prompts=[
+    #     "Task_1.1.1 must precede Task_1.1.2, which in turn should precede Task_1.1.3, ",
+    #     "Task_1.1 must be completed before Task_1.2 starts, and Task_1.2 must be completed before Task_1.3 starts." ,
+    #     "Task_1.1 can be executed independently, after which Task_1.2 can be executed.",
+    #     "Task_1.2.4 must be completed first, followed by Task_1.2.2, then Task_1.2.3, and finally Task_1.2.1.",
+    #     "Task_1.2.4 is always executed first, followed by Task_1.2.3, then Task_1.2.2, and finally Task_1.2.1.",
+    #     "Task_1.2.1 and Task_1.2.2 can be executed independently, and both should eventually be completed.",
+    # ]
+    # for ret in test_prompts:
+    #     print(translater.translate(ret))    
+    #     print('\n','-'*20,'\n')
+    
+    
+    import evaluate
+    import numpy as np
+    # from datasets import load_from_disk
+    from tqdm import tqdm
+
+    # Metric
+    metric = evaluate.load("rouge")
+    tokenized_dataset = load_dataset("json",  data_files={"train":"LTL_datasets/collect/ltl_eng_train_mid_ascii_gptAuged.jsonl","test":"LTL_datasets/collect/ltl_eng_test_mid_ascii_gptAuged.jsonl"})
+    print(tokenized_dataset)
+    # run predictions
+    # this can take ~45 minutes
+    import re 
+    pattern=re.compile("\[Formal LTL\]:\n([\S ]*)\n")
+    predictions, references,input_sentence,output_sentence=[], [] , [], []
+    # with open()
+    for idx in range(len(tokenized_dataset['test']['natural'])):
+        # print(sample)
+        nl=tokenized_dataset['test']['natural'][idx]
+        p = translater.evaluate_model(nl)
+        # print(p,l)
+        input_sentence.append(nl)
+
+        transLTL=pattern.findall(p)
+        # print(p)
+        predictions.append(transLTL[0]) 
+        output_sentence.append(p) 
+        # input_sentence.append(nl)
+        references.append(tokenized_dataset['test']['raw_ltl'][idx])
+        print(idx,'\n',input_sentence[-1],'\nout::\n',output_sentence[-1],'\npre::\n',predictions[-1],'\nref::\n',references[-1],'\n','-'*20,'\n')
+
+    # compute metric
+    rogue = metric.compute(predictions=predictions, references=references, use_stemmer=True)
+
+    # print results
+    print(f"Rogue1: {rogue['rouge1']* 100:2f}%")
+    print(f"rouge2: {rogue['rouge2']* 100:2f}%")
+    print(f"rougeL: {rogue['rougeL']* 100:2f}%")
+    print(f"rougeLsum: {rogue['rougeLsum']* 100:2f}%")
+    eval_output=np.array([input_sentence,predictions,references]).T
+    import pandas as pd 
+    eval_output=pd.DataFrame(eval_output)
+    pd.DataFrame.to_csv(eval_output,"finetuned_model/llama2_13b__mid_asciiaug1"+'/output')
+        # Rogue1: 98.363321%
+        # rouge2: 95.987820%
+        # rougeL: 97.384820%
+        # rougeLsum: 97.382071%
+    exit()
+    flag=True
+    while flag:
+        lines=[""]
+        try:
+            lines.append(input())
+            while True:
+                lines.append(input())
+        except:
+            pass
+        ret ="".join(lines)
+        print(ret)
+        if ret=="":
+            flag=False
+        
+        print(translater.translate(ret))
+
+

finetune/Llama2_13b/llama_test.py

@@ -0,0 +1,138 @@
+# from huggingface_hub import login
+# login()
+from datasets import load_dataset
+from random import randrange
+import torch
+import sys,os
+# sys.path.append('../../../')
+# sys.path.append('../../')
+# sys.path.append('../')
+os.environ['CUDA_VISIBLE_DEVICES'] = "2" 
+# device = torch.device("cuda:0-6" if torch.cuda.is_available() else "cpu")
+from ... import utils as util
+from accelerate import init_empty_weights,infer_auto_device_map,load_checkpoint_in_model,dispatch_model
+# Load dataset from the hub
+# dataset = load_dataset("databricks/databricks-dolly-15k", split="train")
+output_dir = "finetuned_model/"
+datapath='LTL_datasets/collect/'
+exp_name="_mid_ascii"
+base_model_name = "meta-llama/Llama-2-13b-hf"
+dataset = load_dataset("json",  data_files={"train":datapath+"ltl_eng_train"+exp_name+".jsonl","test":datapath+"ltl_eng_test"+exp_name+".jsonl"})
+
+print(dataset)
+# dataset size: 15011
+
+def format_instruction(sample):
+	return f"""### [Instruction]:
+translate natural description in to LTL
+
+### [Input]:
+{sample['natural']}
+
+### [English Response]:
+{util.reAsciiLTL2EngLTL(sample['raw_ltl'])}
+### [Formal Response]:
+{sample['raw_ltl']}
+"""
+
+def preprocess_function(sample,padding="max_length"):
+    # add prefix to the input for t5
+    # print(sample[0])
+    inputs=[
+         f"""### [Instruction]:
+translate natural description in to LTL
+
+### [Input]:
+{sample['natural'][i]}
+
+### [English Response]:
+{util.reAsciiLTL2EngLTL(sample['raw_ltl'][i])}
+### [Formal Response]:
+{sample['raw_ltl'][i]}
+"""
+    for i in (range(len(sample['natural'])))]
+    # inputs = ["## [instruction]: translate natural description in to LTL: ### [natural language]:" + sample['natural'][i]+'### [LTL]:'+sample['raw_ltl'][i] for i in (range(len(sample['natural'])))]
+
+    sample["complete_text"] = inputs
+    return sample
+
+tokenized_dataset = dataset.map(preprocess_function, batched=True)
+print(tokenized_dataset)
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+
+use_flash_attention = False
+
+# Hugging Face model id
+model_id = base_model_name
+# model_id = "meta-llama/Llama-2-7b-hf" # gated
+
+
+output_dir = os.path.join(output_dir, "llama2_13b_"+exp_name+'4')
+
+
+if use_flash_attention:
+    # unpatch flash attention
+    from llama_dp2_patch import unplace_flash_attn_with_attn
+    unplace_flash_attn_with_attn()
+
+import torch
+from peft import AutoPeftModelForCausalLM
+from transformers import AutoTokenizer
+
+# load base LLM model and tokenizer
+model = AutoPeftModelForCausalLM.from_pretrained(
+    output_dir,
+    low_cpu_mem_usage=True,
+    torch_dtype=torch.float16,
+    load_in_4bit=True,
+)
+tokenizer = AutoTokenizer.from_pretrained(output_dir)
+
+
+from datasets import load_dataset
+from random import randrange
+
+
+# Load dataset from the hub and get a sample
+# dataset = load_dataset("databricks/databricks-dolly-15k", split="train")
+# sample = dataset[randrange(len(dataset))]
+
+# prompt = f"""### Instruction:
+# Use the Input below to create an instruction, which could have been used to generate the input using an LLM.
+
+# ### Input:
+# {sample['response']}
+
+# ### Response:
+# """
+
+# input_ids = tokenizer(prompt, return_tensors="pt", truncation=True).input_ids.cuda()
+# # with torch.inference_mode():
+# outputs = model.generate(input_ids=input_ids, max_new_tokens=100, do_sample=True, top_p=0.9,temperature=0.9)
+
+# print(f"Prompt:\n{sample['response']}\n")
+# print(f"Generated instruction:\n{tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0][len(prompt):]}")
+# print(f"Ground truth:\n{sample['instruction']}")
+
+
+
+
+
+from peft import AutoPeftModelForCausalLM
+
+model = AutoPeftModelForCausalLM.from_pretrained(
+    output_dir,
+    low_cpu_mem_usage=True,
+)
+
+# Merge LoRA and base model
+merged_model = model.merge_and_unload()
+
+# Save the merged model
+merged_model.save_pretrained("merged_model",safe_serialization=True)
+tokenizer.save_pretrained("merged_model")
+
+# push merged model to the hub
+# merged_model.push_to_hub("user/repo")
+# tokenizer.push_to_hub("user/repo")

finetune/MIT_NL2TL/NL2TL.py

@@ -0,0 +1,101 @@
+# %%
+from transformers import (AutoModelForSeq2SeqLM, 
+                            AutoTokenizer, 
+                            T5Tokenizer)
+import torch
+import pandas as pd
+from datasets import Dataset, DatasetDict, load_dataset, load_from_disk
+from tqdm import tqdm
+
+#import subprocess
+import sys
+import os
+import argparse
+# from IPython.core import error
+import random
+import numpy as np
+import nltk
+import json
+import csv
+import utils.util as util
+
+# run under conda env minigpt4
+
+class NL2TL():
+	def __init__(self,dirpath='outputdir/') -> None:
+		self.output_dir = dirpath
+		
+		# Here you need to link this path in your Google drive to the place preseving your model weights, e.g., checkpoint-62500
+		# You can download it on the github page
+
+		self.model_checkpoint = "t5-base"
+		self.prefix = "Transform the following sentence into Signal Temporal logic: "
+
+		self.max_input_length = 1024
+		self.max_target_length = 128
+		self.tokenizer = AutoTokenizer.from_pretrained(self.model_checkpoint, model_max_length=self.max_input_length)
+
+		self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+		self.tl_model = AutoModelForSeq2SeqLM.from_pretrained(self.output_dir+"checkpoint-62500").to(self.device)
+
+		# %%
+		import time
+		self.time_start = time.time()
+		self.inputs = [self.prefix + 'At some point (prop_1), and at some point (prop_2), and always do not (prop_4).']
+		self.inputs = self.tokenizer(self.inputs, max_length=self.max_input_length, truncation=True, return_tensors="pt").to(self.device)
+		self.output = self.tl_model.generate(**self.inputs, num_beams=8, do_sample=True, max_length=self.max_target_length)
+		self.decoded_output = self.tokenizer.batch_decode(self.output, skip_special_tokens=True)[0]
+		print(self.decoded_output)
+		self.time_end = time.time()
+		print('Translation time: ', self.time_end-self.time_start)
+		print('\nNL2TL init\n')
+		self.splitJSONfromTXT=util.splitJSONfromTXT
+		# %%
+		# Here are the example test sentences
+		pass
+	def translate(self,inputNLtxt:str=""):
+		inputNLtxt=inputNLtxt.replace("Task_","prop_")
+
+		sentence=inputNLtxt
+		self.inputs = [self.prefix + sentence]
+		self.inputs = self.tokenizer(self.inputs, max_length=self.max_input_length, truncation=True, return_tensors="pt").to(self.device)
+		self.output = self.tl_model.generate(**self.inputs, num_beams=8, do_sample=True, max_length=self.max_target_length)
+		self.decoded_output = self.tokenizer.batch_decode(self.output, skip_special_tokens=True)[0]
+		print('Input sentence: ', sentence)
+		print('Translated STL: ', self.decoded_output)
+		print('\n')
+		            
+		self.decoded_output=self.decoded_output.replace('prop_','Task_')
+		return self.decoded_output
+	def waiting(self):
+		retry=True
+		while retry:
+			inputNL=util.GPTinterface("continue next")
+			if inputNL!="q":
+				Json=self.splitJSONfromTXT(inputNL)
+				print(Json)
+				jsonTree=json.loads("{"+Json[-1]+"}")
+				input_NL=jsonTree["LTL_description"].replace("Task_","prop_")
+				output_TL=self.translate(input_NL)
+				output_TL=output_TL.replace('prop_','Task_')
+				print("\n",output_TL,"\n")
+			else:
+				retry =False
+if __name__=="__main__":
+	# examples=['Stay at (prop_1) for 5 units in the future and stay at (prop_2) for 5 units in the future, and ensure that never (prop_3).',
+	# 'First (prop_1), and then (prop_2), and ensure that never (prop_3).',
+	# 'Start by (prop_1). Then, (prop_2). Lastly, (prop_3).',
+	# 'Guarantee that you (prop_1) and (prop_2)', # Input the natural sentence
+	# '( prop_1 ) and whenever ( prop_2 )',
+	# 'Sooner or later (prop_1)',
+	# 'Repeatedly (prop_1)',
+	# 'At some point, (prop_1).',
+	# 'Do prop_1 but not do prop_2',
+	# 'Do prop_1, do prop_2, do prop_3'] # Input the natural sentence
+	# interface=NL2TL()
+	# for txt in examples:
+	# 	interface.translate(txt)
+
+	interface=NL2TL()
+	interface.waiting()
+

finetune/T5_XXL/t5_lora_evaluate.py

@@ -0,0 +1,95 @@
+import torch
+from peft import PeftModel, PeftConfig
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+# Load peft config for pre-trained checkpoint etc.
+exp_name="_mid_ascii"
+peft_model_id="finetuned_model/results"+exp_name+'2'
+max_target_length=128
+
+config = PeftConfig.from_pretrained(peft_model_id)
+
+# load base LLM model and tokenizer
+model = AutoModelForSeq2SeqLM.from_pretrained(config.base_model_name_or_path,  load_in_8bit=True,  device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path, device_map="auto")
+
+# Load the Lora model
+model = PeftModel.from_pretrained(model, peft_model_id, device_map="auto")
+model.eval()
+
+print("Peft model loaded")
+
+from datasets import load_dataset
+from random import randrange
+
+
+# Load dataset from the hub and get a sample
+datapath='LTL_datasets/collect/'
+dataset = load_dataset("json",  data_files={"train":datapath+"ltl_eng_train"+exp_name+".jsonl","test":datapath+"ltl_eng_test"+exp_name+".jsonl"})
+print(dataset)
+sample = dataset['test'][randrange(len(dataset["test"]))]
+
+input_ids = tokenizer(sample["natural"], return_tensors="pt", truncation=True).input_ids.cuda()
+# with torch.inference_mode():
+outputs = model.generate(input_ids=input_ids, max_new_tokens=max_target_length, do_sample=True, top_p=0.9)
+print(f"input sentence: {sample['natural']}\n{'---'* 20}")
+
+print(f"summary:\n{tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0]}")
+
+
+import evaluate
+import numpy as np
+from datasets import load_from_disk
+from tqdm import tqdm
+
+# Metric
+metric = evaluate.load("rouge")
+
+def evaluate_peft_model(sample,max_target_length=128):
+    # generate summary
+    outputs = model.generate(input_ids=sample["input_ids"].unsqueeze(0).cuda(), do_sample=True, top_p=0.9, max_new_tokens=max_target_length)
+    prediction = tokenizer.decode(outputs[0].detach().cpu().numpy(), skip_special_tokens=True)
+    # decode eval sample
+    # Replace -100 in the labels as we can't decode them.
+    labels = np.where(sample['labels'] != -100, sample['labels'], tokenizer.pad_token_id)
+    # print(labels)
+    labels = tokenizer.decode(labels, skip_special_tokens=True)
+    # print(labels)
+    # Some simple post-processing
+    input_sentence=" ".join(tokenizer.batch_decode(sample["input_ids"].detach().cpu().numpy(), skip_special_tokens=True))
+    print("input sentence: {}\n{}".format(input_sentence,'---'* 20))
+    # output_LTL=tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0]
+    # expect_LTL=labels
+    print(f"pre_LTL:\n{prediction}\nexp_LTL:\n{labels}")
+    return prediction, labels,input_sentence
+
+# load test dataset from distk
+test_dataset = load_from_disk("data/eval"+exp_name+'/').with_format("torch")
+
+# run predictions
+# this can take ~45 minutes
+predictions, references,input_sentence= [] , [], []
+for sample in tqdm(test_dataset):
+    # print(sample)
+    p,l,nl = evaluate_peft_model(sample)
+    # print(p,l)
+    input_sentence.append(nl)
+    predictions.append(p)
+    references.append(l)
+
+# compute metric
+rogue = metric.compute(predictions=predictions, references=references, use_stemmer=True)
+
+# print results
+print(f"Rogue1: {rogue['rouge1']* 100:2f}%")
+print(f"rouge2: {rogue['rouge2']* 100:2f}%")
+print(f"rougeL: {rogue['rougeL']* 100:2f}%")
+print(f"rougeLsum: {rogue['rougeLsum']* 100:2f}%")
+eval_output=np.array([input_sentence,predictions,references]).T
+import pandas as pd 
+eval_output=pd.DataFrame(eval_output)
+pd.DataFrame.to_csv(eval_output,peft_model_id+'/output')
+# Rogue1: 50.386161%
+# rouge2: 24.842412%
+# rougeL: 41.370130%
+# rougeLsum: 41.394230%

finetune/T5_XXL/t5_lora_fintune.py

@@ -0,0 +1,238 @@
+# %%
+from datasets import load_dataset
+import os 
+import sys
+# Load dataset from the hub
+# dataset = load_dataset("samsum")
+# datapath='LTL_datasets/collect/'
+exp_name="/tf-ltl_eng_test_mid_ascii_gptAuged"
+# output_dir = os.path.join(output_dir, "llama2_13b_"+exp_name+'aug1')
+# dataset = load_dataset("json",  data_files={"train":datapath+"ltl_eng_train"+exp_name+".jsonl","test":datapath+"ltl_eng_test"+exp_name+".jsonl"})
+# print(dataset)
+dataset = load_dataset("json",  data_files={"train":"LTL_datasets/collect/ltl_eng_train_mid_ascii_gptAuged.jsonl","test":"LTL_datasets/collect/ltl_eng_test_mid_ascii_gptAuged.jsonl"})
+print(dataset)
+
+os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
+os.environ['CUDA_VISIBLE_DEVICES']='4,5'
+
+
+# %%
+
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+model_id="google/flan-t5-xxl"
+
+# Load tokenizer of FLAN-t5-XL
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+
+# %%
+from datasets import concatenate_datasets
+import numpy as np
+# The maximum total input sequence length after tokenization.
+# Sequences longer than this will be truncated, sequences shorter will be padded.
+tokenized_inputs = concatenate_datasets([dataset["train"], dataset["test"]]).map(lambda x: tokenizer(x["natural"], truncation=True), batched=True, remove_columns=["raw_ltl", "natural"])
+input_lenghts = [len(x) for x in tokenized_inputs["input_ids"]]
+# take 85 percentile of max length for better utilization
+max_source_length = int(np.percentile(input_lenghts, 100))
+print(f"Max source length: {max_source_length}")
+
+# The maximum total sequence length for target text after tokenization.
+# Sequences longer than this will be truncated, sequences shorter will be padded."
+tokenized_targets = concatenate_datasets([dataset["train"], dataset["test"]]).map(lambda x: tokenizer(x["raw_ltl"], truncation=True), batched=True, remove_columns=["raw_ltl", "natural"])
+target_lenghts = [len(x) for x in tokenized_targets["input_ids"]]
+# take 90 percentile of max length for better utilization
+max_target_length = int(np.percentile(target_lenghts, 100))
+print(f"Max target length: {max_target_length}")
+
+
+# %%
+def preprocess_function(sample,padding="max_length"):
+    # add prefix to the input for t5
+    inputs = ["Generate LTL: " + item for item in sample["natural"]]
+
+    # tokenize inputs
+    model_inputs = tokenizer(inputs, max_length=max_source_length, padding=padding, truncation=True)
+
+    # Tokenize targets with the `text_target` keyword argument
+    labels = tokenizer(text_target=sample["raw_ltl"], max_length=max_target_length, padding=padding, truncation=True)
+
+    # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
+    # padding in the loss.
+    if padding == "max_length":
+        labels["input_ids"] = [
+            [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
+        ]
+
+    model_inputs["labels"] = labels["input_ids"]
+    return model_inputs
+
+tokenized_dataset = dataset.map(preprocess_function, batched=True, remove_columns=["natural", "raw_ltl"])
+print(f"Keys of tokenized dataset: {list(tokenized_dataset['train'].features)}")
+
+# save datasets to disk for later easy loading
+tokenized_dataset["train"].save_to_disk("data/train"+exp_name)
+tokenized_dataset["test"].save_to_disk("data/eval"+exp_name)
+
+
+# %%
+from transformers import AutoModelForSeq2SeqLM
+from peft import PeftModel, PeftConfig
+# huggingface hub model id
+model_id = "philschmid/flan-t5-xxl-sharded-fp16"
+model = AutoModelForSeq2SeqLM.from_pretrained(model_id, load_in_8bit=True, device_map="auto")
+
+
+
+# peft_model_id="finetuned_model/results"+"_mid_ascii"
+# config = PeftConfig.from_pretrained(peft_model_id)
+# # load base LLM model and tokenizer
+# model = AutoModelForSeq2SeqLM.from_pretrained(config.base_model_name_or_path,  load_in_8bit=True,  device_map="auto")
+# tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path, device_map="auto")
+
+# # Load the Lora model
+# model = PeftModel.from_pretrained(model, peft_model_id, device_map="auto")
+# # load model from the hub
+
+print(model)
+# exit()
+
+# %%
+from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training, TaskType
+
+# Define LoRA Config
+lora_config = LoraConfig(
+ r=16,
+ lora_alpha=32,
+ target_modules=["q", "v"],
+ lora_dropout=0.05,
+ bias="none",
+ task_type=TaskType.SEQ_2_SEQ_LM
+)
+# prepare int-8 model for training
+model = prepare_model_for_int8_training(model)
+
+# add LoRA adaptor
+model = get_peft_model(model, lora_config)
+model.print_trainable_parameters()
+
+# trainable params: 18874368 || all params: 11154206720 || trainable%: 0.16921300163961817
+
+
+# %%
+from transformers import DataCollatorForSeq2Seq
+
+# we want to ignore tokenizer pad token in the loss
+label_pad_token_id = -100
+# Data collator
+data_collator = DataCollatorForSeq2Seq(
+    tokenizer,
+    model=model,
+    label_pad_token_id=label_pad_token_id,
+    pad_to_multiple_of=8
+)
+
+
+# %%
+from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments
+
+output_dir="lora-flan-t5-xxl"
+
+# Define training args
+training_args = Seq2SeqTrainingArguments(
+    output_dir=output_dir,
+	auto_find_batch_size=True,
+    learning_rate=1e-3, # higher learning rate
+    num_train_epochs=5,
+    logging_dir=f"{output_dir}/logs",
+    logging_strategy="steps",
+    logging_steps=500,
+    save_strategy="no",
+    report_to="tensorboard",
+)
+
+# Create Trainer instance
+trainer = Seq2SeqTrainer(
+    model=model,
+    args=training_args,
+    data_collator=data_collator,
+    train_dataset=tokenized_dataset["train"],
+)
+model.config.use_cache = False  # silence the warnings. Please re-enable for inference!
+
+
+# %%
+# train model
+trainer.train()
+
+
+# %%
+# Save our LoRA model & tokenizer results
+peft_model_id="finetuned_model/"+exp_name
+trainer.model.save_pretrained(peft_model_id)
+tokenizer.save_pretrained(peft_model_id)
+# if you want to save the base model to call
+# trainer.model.base_model.save_pretrained(peft_model_id)
+    
+
+
+
+import evaluate
+import numpy as np
+from datasets import load_from_disk
+from tqdm import tqdm
+
+# Metric
+metric = evaluate.load("rouge")
+
+def evaluate_peft_model(sample,max_target_length=128):
+    # generate summary
+    outputs = model.generate(input_ids=sample["input_ids"].unsqueeze(0).cuda(), do_sample=True, top_p=0.9, max_new_tokens=max_target_length)
+    prediction = tokenizer.decode(outputs[0].detach().cpu().numpy(), skip_special_tokens=True)
+    # decode eval sample
+    # Replace -100 in the labels as we can't decode them.
+    labels = np.where(sample['labels'] != -100, sample['labels'], tokenizer.pad_token_id)
+    # print(labels)
+    labels = tokenizer.decode(labels, skip_special_tokens=True)
+    # print(labels)
+    # Some simple post-processing
+    input_sentence=" ".join(tokenizer.batch_decode(sample["input_ids"].detach().cpu().numpy(), skip_special_tokens=True))
+    # print("input sentence: {}\n{}".format(input_sentence,'---'* 20))
+    # output_LTL=tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0]
+    # expect_LTL=labels
+    # print(f"pre_LTL:\n{prediction}\nexp_LTL:\n{labels}")
+    return prediction, labels,input_sentence
+
+# load test dataset from distk
+test_dataset = load_from_disk("data/eval"+exp_name+'/').with_format("torch")
+
+# run predictions
+# this can take ~45 minutes
+predictions, references,input_sentence= [] , [], []
+idx=0
+for sample in tqdm(test_dataset):
+    # print(sample)
+    p,l,nl = evaluate_peft_model(sample)
+    # print(p,l)
+    input_sentence.append(nl)
+    predictions.append(p)
+    references.append(l)
+    idx+=1
+    print(idx,'\n',input_sentence[-1],'\npre::\n',predictions[-1],'\nref::\n',references[-1],'\n','-'*20,'\n')
+
+# compute metric
+rogue = metric.compute(predictions=predictions, references=references, use_stemmer=True)
+
+# print results
+print(f"Rogue1: {rogue['rouge1']* 100:2f}%")
+print(f"rouge2: {rogue['rouge2']* 100:2f}%")
+print(f"rougeL: {rogue['rougeL']* 100:2f}%")
+print(f"rougeLsum: {rogue['rougeLsum']* 100:2f}%")
+eval_output=np.array([input_sentence,predictions,references]).T
+import pandas as pd 
+eval_output=pd.DataFrame(eval_output)
+pd.DataFrame.to_csv(eval_output,peft_model_id+'/output')
+# Rogue1: 98.292692%
+# rouge2: 95.766211%
+# rougeL: 97.086188%
+# rougeLsum: 97.084262%

finetune/T5_XXL/t5_realtime_evaluate.py

@@ -0,0 +1,69 @@
+import torch
+from peft import PeftModel, PeftConfig
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import sys
+# sys.path.append("..")
+# sys.path.append("../../")
+from ... import utils as util
+# Load peft config for pre-trained checkpoint etc.
+
+class T5XXL_NL2TL_translator():
+    def __init__(self) -> None:
+        # exp_name="_mid_ascii"
+        peft_model_id="model_weight/tf-ltl_eng_test_mid_ascii_gptAuged"
+        self.max_target_length=128
+
+        self.config = PeftConfig.from_pretrained(peft_model_id)
+
+        # load base LLM model and tokenizer
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(self.config.base_model_name_or_path,  load_in_8bit=True,  device_map="auto")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.config.base_model_name_or_path, device_map="auto")
+
+        # Load the Lora model
+        self.model = PeftModel.from_pretrained(self.model, peft_model_id, device_map="auto")
+        self.model.eval()
+
+        print("Peft model loaded")
+
+        pass
+    def translate(self,input:str=""):
+        input_prompt= "Generate LTL: " + input
+        replace=util.Task2Preplacer()
+        
+        input_prompt=replace.reTask2P(input_prompt)
+        # print(predicter( replace.reTask2P(input_prompt)))
+        print(input_prompt)
+        input_ids = self.tokenizer(input_prompt, return_tensors="pt", truncation=True).input_ids.cuda()
+        outputs = self.model.generate(input_ids=input_ids, max_new_tokens=self.max_target_length, do_sample=True, top_p=0.9)
+        output_txt= self.tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0]
+        print(output_txt)
+        return replace.reP2Task(output_txt)
+
+
+if __name__=="__main__":
+    test_prompts=[
+        "Task_1.1.1 must precede Task_1.1.2, which in turn should precede Task_1.1.3, ensuring that arranging fruits happens before preparing vegetables and prepping eggs and meats is done last.",
+        "Task_1.1 must be completed before Task_1.2 starts, and Task_1.2 must be completed before Task_1.3 starts." 
+    ]
+
+    translater=T5XXL_NL2TL_translator()
+
+    for ret in test_prompts:
+        print(translater.translate(ret))    
+    
+    flag=True
+    while flag:
+        lines=[""]
+        try:
+            lines.append(input())
+            while True:
+                lines.append(input())
+        except:
+            pass
+        ret ="".join(lines)
+        print(ret)
+        if ret=="":
+            flag=False
+        
+        print(translater.translate(ret))
+

finetune/data_augmentation/GPTbasedAug.py

@@ -0,0 +1,100 @@
+import json 
+import re
+import sys,os 
+import numpy as np 
+# sys.path.append(os.path.join(os.path.dirname(__file__), '../../'))
+from ... import utils as util
+import time
+l,r=40,2000
+# has been added: 0,40 
+# the range for data to augmentation 
+batch_size=20
+# number of re describe codes
+
+
+dataPath="LTL_datasets/collect/"
+data_eng_path=os.path.join(dataPath,"Cleaned_ENG.txt")
+data_ltl_path=os.path.join(dataPath,"Cleaned_LTL.txt")
+
+data_eng_output_path=os.path.join(dataPath,"eng_gpt_auged2.txt")
+data_ltl_output_path=os.path.join(dataPath,"ltl_mid_order_ascii_gpt_auged2.txt")
+data_src_output_path=os.path.join(dataPath,"idxsrc_gpt_auged2.txt")
+
+with open(data_ltl_path) as txt:
+    content = txt.readlines() 
+    txt.close()
+ltl =np.array(content)
+with open(data_eng_path) as txt:
+    content = txt.readlines() 
+    txt.close()
+eng =np.array(content)
+
+print(len(ltl))
+
+GPTinterface=util.GPTinterface(JSONlog=True,exp_PATH=dataPath)
+
+import random 
+np.random.seed(42)
+
+idx=np.arange(len(ltl))
+np.random.shuffle(idx)
+
+messages=[
+        {
+        "role": "system",
+        "content": """1. Herer are some one sentence examples in a way that is normally used to interpret the safe or co-safe property in linear temporal logic, please remember and imitate the language style in the examples below 
+P02 and P03 can occur independently and either may be executed without affecting the other.
+P07 must precede P17, which in turn should precede P15, ensuring that P07 happens before P17 and P15 is done last.
+Always (P08 precedes P09) and Eventually (P08 is executed) and Eventually (P09 is executed).
+Globally, P02 should be completed before P03 eventually starts.
+Eventually, P08 and P09 should both be completed, and they can be done in any order.
+P02 and P04 can be executed concurrently, while P03 can only be executed once P02 has been completed.
+P07 must be completed before P17, P15, and P02 can be started. P17 must be completed before P15 and P02 can be started. P15 must be completedbefore P02 can be started."
+P08 is a prerequisite for P09, P09 is a prerequisite for P10, and P10 is a prerequisite for P16.
+P06 must be completed before P14, P14 must be completed before P11, and P11 must be completed before P12.
+P02 and P05 are always possible to be executed, while the possibility of executing P03 and P04 is contingent upon the completion of P02.
+P07 must be executed, and only after P07 is completed can P17 be executed, and only after P17 is completed can P15 be executed.
+P08 can be executed independently, after which P09 can be executed.
+P06 must be completed first, followed by P14, then P11, and finally P12.
+P19 is always executed first, followed by P13, then P18, and finally P05.
+P15 and P14 can be executed independently, and both should eventually be completed.
+P07 must be completed before P17, and P17 must be completed before P15
+P06 must be completed before P14 begins, and P14 must be completed before P11 begins"""
+        },{
+            "role":"user",
+            "content":"first go to P01 and then go to P20, always avoiding P02"
+        }
+    ]
+input_content="B. re describe this instruction using the style above\n"
+input_LTL=""
+input_idx=""
+count=0
+pattern=re.compile("[0-9]{2}\. ([\S ]*)\n")
+
+for i in range(l,r):
+    if count>=batch_size:
+        count=0
+        messages[1]["content"]=input_content
+        GPTreturn=GPTinterface.communicate(messages=messages)
+        reDescription=pattern.findall(GPTreturn+'\n')
+        if len(reDescription)==batch_size:
+            with open(data_eng_output_path ,"a") as f:
+                for j in reDescription:
+                    f.write(j)
+                    f.write('\n')
+            with open(data_ltl_output_path,"a") as f:
+                f.write(input_LTL)
+            with open(data_src_output_path,"a") as f:
+                f.write(input_idx)
+        input_content="B. re describe this instruction using the style above\n"
+        input_LTL=""
+        input_idx=""
+        time.sleep(np.random.random()*5)
+    else:
+        count+=1
+        input_content+="{:0>2d}. {}".format(count,eng[idx[i]])
+        input_LTL+="{}".format(ltl[idx[i]])
+        input_idx+="{}\n".format(idx[i])
+    
+
+    

finetune/data_augmentation/dataset_creator.py

@@ -0,0 +1,76 @@
+import json 
+import numpy as np 
+import os,sys
+# sys.path.append(os.path.join(os.path.dirname(__file__), '../../'))
+from ... import utils as util
+class DataPreprocess():
+    def __init__(self,data_path="LTL_datasets/collect") -> None:
+        self.data_path=data_path
+        self.train_valid_split=0.1
+        pass
+    def txtdataReader(self):
+        LTL_list=[
+            # 'ltl_mid_order_ascii.txt',
+            # 'ltl_mid_order_ascii_gpt_auged.txt',
+            # 'ltl_mid_order_ascii_gpt_auged2.txt',
+            'Cleaned_LTL.txt'
+        ]
+        ENG_list=[
+            # 'eng.txt',
+            # 'eng_gpt_auged.txt',
+            # 'eng_gpt_auged2.txt'
+            'Cleaned_ENG.txt'
+        ]
+        content=[]
+        for filename in LTL_list:
+            with open(os.path.join(self.data_path,filename)) as txt:
+                content += txt.readlines() 
+                txt.close()
+        self.ltl =np.array(content)
+
+        content=[]
+        for filename in ENG_list:
+            with open(os.path.join(self.data_path,filename)) as txt:
+                content += txt.readlines() 
+                txt.close()
+        self.eng =np.array(content)
+        print(len(self.ltl))
+
+    def JSONdataCreate(self):
+        self.txtdataReader()
+        self.JSONWriter()
+
+    def JSONWriter(self):
+        np.random.seed(42)
+        # idx=np.random.shuffle( np.arange(len(ltl)))
+        self.idx=np.arange(len(self.ltl))
+        np.random.shuffle(self.idx)
+        with open(self.data_path+"/ltl_eng_train_mid_ascii_gptAuged.jsonl","w") as f:
+            for i in range(int(len(self.ltl)*(1-self.train_valid_split))):
+                json.dump({"natural":self.eng[self.idx[i]],"raw_ltl":self.ltl[self.idx[i]],"id":str(self.idx[i])},f)
+                f.write('\n')
+        with open(self.data_path+"/ltl_eng_test_mid_ascii_gptAuged.jsonl","w") as f:
+            for i in range(int(len(self.ltl)*(1-self.train_valid_split)),len(self.ltl)):
+                json.dump({"natural":self.eng[self.idx[i]],"raw_ltl":self.ltl[self.idx[i]],"id":str(self.idx[i])},f)
+                f.write('\n')
+
+    def dataCheck(self):
+        self.txtdataReader()
+        checker=util.LTLChecker()
+        with open(os.path.join(self.data_path,"Cleaned_LTL.txt"),"a") as passed_LTL:
+            with open(os.path.join(self.data_path,"Cleaned_ENG.txt"),"a") as passed_ENG:
+                with open(os.path.join(self.data_path,"UNCleaned_num.txt"),"a") as unpassed_row:
+                    with open(os.path.join(self.data_path,"UNCleaned_LTL.txt"),"a") as unpassed_LTL:
+                        with open(os.path.join(self.data_path,"UNCleaned_ENG.txt"),"a") as unpassed_ENG:
+                            for id in range(len(self.ltl)):
+                                if checker.AP_CorrCheck(self.ltl[id],self.eng[id]):
+                                    passed_LTL.write(self.ltl[id])
+                                    passed_ENG.write(self.eng[id])
+                                else:
+                                    unpassed_row.write("{}\n".format(id))
+                                    unpassed_LTL.write(self.ltl[id])
+                                    unpassed_ENG.write(self.eng[id])
+                            
+if __name__=="__main__":
+    # DataPreprocess().dataCheck()
+    DataPreprocess().JSONdataCreate()

finetune/mistral7b/finetune.py

@@ -0,0 +1,353 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+# device = "cuda" # the device to load the model onto
+# from huggingface_hub import login
+# login()
+import json
+import numpy as np
+import sys,os
+from datasets import load_dataset
+import torch
+from transformers import (AutoModelForCausalLM,
+                         AutoTokenizer,
+                         BitsAndBytesConfig,
+                         TrainingArguments,
+                         pipeline,
+                         logging,
+                         TrainerCallback)
+from peft import LoraConfig, PeftConfig, prepare_model_for_kbit_training, get_peft_model
+from trl import SFTTrainer
+from accelerate import infer_auto_device_map,init_empty_weights
+import wandb
+from datasets import concatenate_datasets
+import numpy as np
+# sys.path.append('../../../')
+# sys.path.append('../../')
+# sys.path.append('../')
+# os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+# os.environ['CUDA_VISIBLE_DEVICES'] = "5,6,7" 
+# device = torch.device("cuda:0-6" if torch.cuda.is_available() else "cpu")
+sys.path.append(os.path.join(os.path.dirname(__file__), '../../'))
+# import utils.util as util
+# Load dataset from the hub
+# dataset = load_dataset("samsum")
+device='cuda'
+np.random.seed(42)
+output_dir = "/home/user/xsj/model_weight/"
+datapath='/home/user/xsj/NL2TL-dataset/collect2'
+exp_name="_mid_ascii_0327_eos_2"
+explainer_files=['LTLexplain_0.json','LTLexplain_1.json','LTLexplain_2.json','LTLexplain_3.json']
+explainer_dic={}
+for path in explainer_files:
+    with open(os.path.join(datapath,path)) as f:
+        LTLlist=json.load(f)
+    for key in LTLlist.keys():
+        if isinstance(LTLlist[key],dict):
+            if not (key in explainer_dic):
+                explainer_dic[key]=[]
+            explainer_dic[key].append(LTLlist[key]['translate'])
+            sp=LTLlist[key]['explain'].split("means that")
+            if len(sp)>1:
+                explainer_dic[key].append(sp[1])
+    
+base_model_name = "mistralai/Mistral-7B-Instruct-v0.2"
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit = True,
+    bnb_4bit_use_double_quant = False,
+    bnb_4bit_quant_type = 'nf4',
+    bnb_4bit_compute_dtype = getattr(torch, "float16")
+)
+bnb_config = BitsAndBytesConfig(
+    load_in_8bit = True,
+    # llm_int8_threshold=200.0
+    # bnb_4bit_use_double_quant = False,
+    # bnb_4bit_quant_type = 'nf4',
+    # bnb_4bit_compute_dtype = getattr(torch, "float16")
+)
+import os
+os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
+# os.environ['CUDA_VISIBLE_DEVICES']='0'
+device_map="auto"
+# torch.cuda.set_device(7)
+# device_map={'':torch.cuda.current_device()}
+# device_map = {'':'cuda:7'}
+# model_dir为模型的路径或名称
+# config = AutoConfig.from_pretrained(base_model_name, trust_remote_code=True)
+# with init_empty_weights():
+#     base_model = AutoModelForCausalLM.from_pretrained(
+#     base_model_name,
+#     from_tf=bool(".ckpt" in base_model_name),
+#     quantization_config=bnb_config,
+#     device_map=device_map,
+#     trust_remote_code=True,
+#     use_auth_token=True
+#     )
+
+# map_list = {5:"15GB", 6:"15GB",7:"15GB"}    # 对应不同卡号限制的内存量
+# map_list = {7:"15GB",}    # 对应不同卡号限制的内存量
+# no_split_modules = base_model._no_split_modules
+# device_map = infer_auto_device_map(base_model, max_memory=map_list, no_split_module_classes=no_split_modules)
+
+
+dataset = load_dataset("json",  data_files={"train":os.path.join(datapath,"ltl_eng_train_mid_ascii_gptAuged.jsonl"),"test":os.path.join(datapath,"ltl_eng_test_mid_ascii_gptAuged.jsonl")})
+print(dataset)
+
+
+
+# tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+# , add_eos_token=True,trust_remote_code=True)
+# NOTE no one says whether the add eos token need to be added, but if we do not add this, the generate will continue until reach the max_new_tokens
+# when add add_eos_token, it always failed 
+# if use this it will generate somthing other
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = 'right'
+# print(tokenizer.eos_token_id)
+# 2
+# print(tokenizer.bos_token_id)
+# 1
+# print(tokenizer._convert_token_to_id(tokenizer.bos_token))
+
+def preprocess_function(sample,padding="max_length"):
+    # add prefix to the input for t5
+    # print(sample[0])
+    inputs=[
+         f"""### Instruction:
+translate natural description to linear temproal logic, first translate into a logical way, and then translate into linear temproal logic, pay specific attention to brackets '()'
+
+### Natural Language Task:
+{sample['natural'][i].strip()}
+
+### Logic Translation:
+{explainer_dic[sample['raw_ltl'][i].strip()][np.random.randint(0,len(explainer_dic[sample['raw_ltl'][i].strip()]))]}
+
+### linear temproal logic:
+{sample['raw_ltl'][i].strip()}
+</s>""".lower()
+    # NOTE it seems the eos is needed, the bos is not needed(the bos will be automatically added) 
+    for i in (range(len(sample['natural'])))]
+    # inputs = ["## [instruction]: translate natural description in to LTL: ### [natural language]:" + sample['natural'][i]+'### [LTL]:'+sample['raw_ltl'][i] for i in (range(len(sample['natural'])))]
+
+    sample["complete_text"] = inputs
+    return sample
+# method1
+# tokenized_dataset = dataset.map(preprocess_function, batched=True)
+# method2
+def preprocess_function2(sample,padding="max_length"):
+    # add prefix to the input for t5
+    # print(sample[0])
+    inputs=[
+         tokenizer.apply_chat_template(
+    [
+        {"role": "user", "content": "translate natural description to linear temproal logic, first translate into a logical expression, and then translate into linear temproal logic, please pay specific attention to logic grammar, the natural language task is {}".format(sample['natural'][i].strip())},
+        {"role": "assistant", "content": "logic expression is {}, and LTL is {} .".format(
+            explainer_dic[sample['raw_ltl'][i].strip()][np.random.randint(0,len(explainer_dic[sample['raw_ltl'][i].strip()]))],
+            sample['raw_ltl'][i].strip()
+            )
+        },
+        # {"role": "user", "content": " pay specific attention to brackets '()', linear temproal logic is"},
+        # {"role": "assistant", "content": "LTL is {} .".format(
+        #     sample['raw_ltl'][i].strip()
+        #     )
+        # }
+    ],tokenize=False)
+    # NOTE it seems the eos is needed, the bos is not needed(the bos will be automatically added) 
+    for i in (range(len(sample['natural'])))]
+    # inputs = ["## [instruction]: translate natural description in to LTL: ### [natural language]:" + sample['natural'][i]+'### [LTL]:'+sample['raw_ltl'][i] for i in (range(len(sample['natural'])))]
+
+    sample["complete_text"] = inputs
+    return sample
+tokenized_dataset = dataset.map(preprocess_function2, batched=True)
+print(f"Keys of tokenized dataset: {list(tokenized_dataset['train'].features)}")
+
+# save datasets to disk for later easy loading
+# tokenized_dataset["train"].save_to_disk("data/train"+exp_name)
+# tokenized_dataset["test"].save_to_disk("data/eval"+exp_name)
+
+class PeftSavingCallback(TrainerCallback):
+    def on_save(self, args, state, control, **kwargs):
+        checkpoint_path = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")
+        kwargs["model"].save_pretrained(checkpoint_path)
+
+        if "pytorch_model.bin" in os.listdir(checkpoint_path):
+            os.remove(os.path.join(checkpoint_path, "pytorch_model.bin"))
+callbacks = [PeftSavingCallback]
+
+peft_config = LoraConfig(
+    lora_alpha=16,
+    lora_dropout=0.05,
+    r=128,
+    bias="none",
+    task_type="CAUSAL_LM",
+    target_modules=["q_proj", "v_proj"]
+)
+
+
+training_arguments = TrainingArguments(
+    output_dir=output_dir,
+    logging_dir = os.path.join(output_dir,"logs"),
+    per_device_train_batch_size=1,
+    num_train_epochs=3,
+    gradient_accumulation_steps=8,
+    optim="paged_adamw_32bit",
+    save_strategy='epoch',
+    logging_steps=25,
+    learning_rate=2e-4,
+    weight_decay=0.001,
+    fp16=True,
+    bf16=False,
+    max_grad_norm=0.3,
+    max_steps=-1,
+    warmup_ratio = 0.05,
+    group_by_length=True,
+    lr_scheduler_type="cosine",
+    report_to="wandb",
+    evaluation_strategy="epoch",
+    do_eval=True,
+    run_name = base_model_name+exp_name,
+    disable_tqdm=False
+)
+import os
+output_dir = os.path.join(output_dir, "mistral7b"+exp_name+'aug1_quat8')
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
+    from_tf=bool(".ckpt" in base_model_name),
+    quantization_config=bnb_config,
+    device_map=device_map,
+    trust_remote_code=True,
+    use_auth_token=True
+)
+base_model.config.use_cache = False
+
+# More info: https://github.com/huggingface/transformers/pull/24906
+base_model.config.pretraining_tp = 1 
+
+base_model.gradient_checkpointing_enable()
+base_model = prepare_model_for_kbit_training(base_model)
+base_model = get_peft_model(base_model, peft_config)
+
+trainer = SFTTrainer(
+    model=base_model,
+    train_dataset=tokenized_dataset['train'],
+    eval_dataset=tokenized_dataset['test'],
+    peft_config=peft_config,
+    dataset_text_field="complete_text",
+    max_seq_length=512,
+    tokenizer=tokenizer,
+    args=training_arguments,
+    callbacks=callbacks,
+    packing=False,
+)
+wandb.login()
+trainer.train()
+trainer.model.save_pretrained(output_dir)
+# trainer.model.save_pretrained(output_dir)
+tokenizer.save_pretrained(output_dir)
+
+wandb.finish()
+
+# check
+print('model dir',output_dir)
+from peft import AutoPeftModelForCausalLM
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = AutoPeftModelForCausalLM.from_pretrained(output_dir,
+    from_tf=bool(".ckpt" in output_dir),
+    quantization_config=bnb_config,
+    device_map=device_map,
+    trust_remote_code=True,
+    use_auth_token=True
+)
+tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
+print(tokenizer.default_chat_template)
+def evaluate_model(input_text):
+    input_text =f"""### Instruction:
+translate natural description to linear temproal logic, first translate into a logical way, and then translate into linear temproal logic, pay specific attention to brackets '()' ### Natural Language Task:
+{input_text}""".lower()
+    inputs = tokenizer(input_text, return_tensors="pt").to(device)
+    print(inputs)
+    outputs = model.generate(input_ids=inputs["input_ids"].to("cuda"), attention_mask=inputs["attention_mask"].to("cuda"), max_new_tokens=512, pad_token_id=tokenizer.eos_token_id)
+
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+def evaluate_model2(input_text):
+    messages=[
+        {"role": "user", "content": "translate natural description to linear temproal logic, first translate into a logical way, and then translate into linear temproal logic, pay specific attention to brackets '()', natural language task: {}".format(input_text)},
+    ]
+
+    encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt").to(device)
+    outputs = model.generate(encodeds, max_new_tokens=512)
+    # , pad_token_id=tokenizer.eos_token_id)
+#     input_text =f"""### Instruction:
+# translate natural description to linear temproal logic, first translate into a logical way, and then translate into linear temproal logic, pay specific attention to brackets '()' ### Natural Language Task:
+# {input_text}""".lower()
+#     inputs = tokenizer(input_text, return_tensors="pt").to(device)
+#     print(inputs)
+#     outputs = model.generate(input_ids=inputs["input_ids"].to("cuda"), attention_mask=inputs["attention_mask"].to("cuda"), max_new_tokens=512, pad_token_id=tokenizer.eos_token_id)
+
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+# if __name__=='__main__':
+import evaluate
+import numpy as np
+from datasets import load_from_disk
+from tqdm import tqdm
+
+# Metric
+metric = evaluate.load("rouge")
+
+
+# load test dataset from distk
+# test_dataset = load_from_disk("data/eval"+exp_name+'/').with_format("torch")
+
+# run predictions
+# this can take ~45 minutes
+import re 
+pattern=re.compile("linear temproal logic is ([\S ]*)")
+predictions, references,input_sentence,output_sentence=[], [] , [], []
+for idx in range(len(tokenized_dataset['test']['natural'])):
+    # print(sample)
+    nl=tokenized_dataset['test']['natural'][idx]
+    p = evaluate_model2(nl)
+    # print(p,l)
+    input_sentence.append(nl)
+
+    transLTL=pattern.findall(p)
+    print(p)
+    if transLTL[0][-1]=='.':
+        transLTL[0]=transLTL[0][:-1].strip()
+    else:
+        transLTL[0]=transLTL[0].strip()
+    predictions.append(transLTL[0]) 
+    output_sentence.append(p) 
+    input_sentence.append(p)
+    references.append(tokenized_dataset['test']['raw_ltl'][idx].strip())
+    print(input_sentence[-1],'\nout::\n',output_sentence[-1],'\npre::\n',predictions[-1],'\nref::\n',references[-1],'\n','-'*20,'\n')
+
+# compute metric
+rogue = metric.compute(predictions=predictions, references=references, use_stemmer=True)
+
+# print results
+print(f"Rogue1: {rogue['rouge1']* 100:2f}%")
+print(f"rouge2: {rogue['rouge2']* 100:2f}%")
+print(f"rougeL: {rogue['rougeL']* 100:2f}%")
+print(f"rougeLsum: {rogue['rougeLsum']* 100:2f}%")
+eval_output=np.array([input_sentence,predictions,references]).T
+import pandas as pd 
+eval_output=pd.DataFrame(eval_output)
+pd.DataFrame.to_csv(eval_output,output_dir+'/output')
+
+exit()
+messages = [
+    {"role": "user", "content": "What is your favourite condiment?"},
+    {"role": "assistant", "content": "Well, I'm quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever I'm cooking up in the kitchen!"},
+    {"role": "user", "content": "Do you have mayonnaise recipes?"}
+]
+
+encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt")
+
+model_inputs = encodeds.to(device)
+model.to(device)
+
+generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
+decoded = tokenizer.batch_decode(generated_ids)
+print(decoded[0])

finetune/mistral7b/prediction.py

@@ -0,0 +1,308 @@
+# from huggingface_hub import login
+# login()
+import sys,os
+from datasets import load_dataset
+import torch
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizer, TrainingArguments
+# from peft import LoraConfig
+# from trl import SFTTrainer
+# from accelerate import infer_auto_device_map,init_empty_weights
+
+# sys.path.append(os.path.join(os.path.dirname(__file__), '../../'))
+from NL2HLTLtaskPlanner.utils import Task2Preplacer
+from NL2HLTLtaskPlanner.utils import LTLChecker
+import re 
+from datasets import concatenate_datasets
+import numpy as np
+from peft import AutoPeftModelForCausalLM
+os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
+# os.environ['CUDA_VISIBLE_DEVICES']='3'
+
+
+
+class Mistral_NL2TL_translator():
+    def __init__(self,
+                 output_dir = "/home/user/xsj/model_weight",  
+                 tuned_model_name="mistral7b_mid_ascii_0327_eos_2aug1_quat8",
+                #  CUDA_device='0',
+                 quat=True,
+                 replacer=Task2Preplacer) -> None:
+        # os.environ['CUDA_VISIBLE_DEVICES']=CUDA_device
+        self.device_map="auto"
+        self.model_dir = os.path.join(output_dir, tuned_model_name)
+        # check
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        # AutoPeftModelForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf")
+
+
+        # quantconfig = BitsAndBytesConfig(
+        #     load_in_8bit=True,
+        #     bnb_8bit_quant_type="nf4",
+        #     bnb_8bit_use_double_quant=True,
+        #     bnb_8bit_compute_dtype=torch.bfloat16,
+        # )
+        # if quat==False:
+        #     self.model = AutoPeftModelForCausalLM.from_pretrained(self.output_dir, device_map=self.device_map, torch_dtype=torch.bfloat16)
+        #     # ICL super man可以不量化
+        # else:
+        #     self.model = AutoPeftModelForCausalLM.from_pretrained(self.output_dir,device_map=self.device_map,  torch_dtype=torch.float16,
+        #     load_in_8bit=True)
+        #     # quantization_config=quantconfig)
+        self.bnb_config = BitsAndBytesConfig(
+            load_in_4bit = True,
+            bnb_4bit_use_double_quant = False,
+            bnb_4bit_quant_type = 'nf4',
+            bnb_4bit_compute_dtype = getattr(torch, "float16")
+        )
+        self.bnb_config = BitsAndBytesConfig(
+            load_in_8bit = True,
+            # llm_int8_threshold=200.0
+            # bnb_4bit_use_double_quant = False,
+            # bnb_4bit_quant_type = 'nf4',
+            # bnb_4bit_compute_dtype = getattr(torch, "float16")
+        )
+        # self.bnb_config = BitsAndBytesConfig(
+        #     load_in_8bit = False,
+        #     load_in_4bit = False,
+        #     # llm_int8_threshold=200.0
+        #     # bnb_4bit_use_double_quant = False,
+        #     # bnb_4bit_quant_type = 'nf4',
+        #     # bnb_4bit_compute_dtype = getattr(torch, "float16")
+        # )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_dir,
+            from_tf=bool(".ckpt" in self.model_dir),
+            quantization_config=self.bnb_config,
+            device_map=self.device_map,
+            trust_remote_code=True,
+            use_auth_token=True
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir)
+        # , trust_remote_code=True,add_eos_token=True,)
+        # tokenizer = AutoTokenizer.from_pretrained(base_model_name, add_eos_token=True,trust_remote_code=True)
+        # NOTE no one says whether the add eos token need to be added, but if we do not add this, the generate will continue until reach the max_new_tokens, 
+        # when in predict model, do not use the add_eos_token=True, as the tokenizer will automatically add <\s> to the input, and thus the output will be inregular
+        # when add add_eos_token, it always failed 
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.padding_side = 'right'
+        print(self.tokenizer.eos_token_id)
+        # 2
+        print(self.tokenizer.bos_token_id)
+        # 1
+        # print(tokenizer._convert_token_to_id(tokenizer.bos_token))
+
+        print("NL2TL model loaded")
+        
+        self.replacer=replacer
+        self.ltlChecker=LTLChecker()
+        pass
+
+        # print('NL2TL llama translate test:')
+        # self.translate("Task_1.1 must be done, and Task_1.2 should be finished before Task_1.1")
+    def evaluate_model(self, input_text):
+        self.pattern=re.compile("linear temproal logic is ([\S ]*).")
+        messages=[
+            {"role": "user", "content": "translate natural description to linear temproal logic, first translate into a logical way, and then translate into linear temproal logic, pay specific attention to brackets '()', natural language task: {}".format(input_text.strip())},
+        ]
+
+        encodeds = self.tokenizer.apply_chat_template(messages, return_tensors="pt").to(self.device)
+        outputs = self.model.generate(encodeds, max_new_tokens=512, pad_token_id=self.tokenizer.eos_token_id)
+
+        p=self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        print('model output:',p)
+        transLTL=self.pattern.findall(p)[0]
+        if transLTL[-1]=='.':
+            transLTL=transLTL[:-1].strip()
+        else:
+            transLTL=transLTL.strip()
+        transLTL=self.ltlChecker.right_barkets_remover(transLTL)
+        print('transLTL:\n',transLTL)
+        return transLTL
+    def evaluate_model2(self, input_text):
+        self.pattern=re.compile("LTL is ([\S ]*).")
+        messages=[
+            {"role": "user", "content": "translate natural description to linear temproal logic, first translate into a logical expression, and then translate into linear temproal logic, the natural language task is {}".format(input_text.strip())},
+        ]
+        encodeds = self.tokenizer.apply_chat_template(messages, return_tensors="pt").to(self.device)
+        outputs = self.model.generate(encodeds, max_new_tokens=512, pad_token_id=self.tokenizer.eos_token_id)
+        p=self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        print('---model output 1:\n',p)
+        # messages=[
+        #     {"role": "user", "content": "translate natural description to linear temproal logic, first translate into a logical expression, and then translate into linear temproal logic, the natural language task is {}".format(input_text.strip())},
+        #     {"role": "assistant", "content":p
+        #     },
+        #     {"role": "user", "content": " pay specific attention to brackets '()', given your linear temproal logic translation"},
+        # ]
+        
+        # encodeds = self.tokenizer.apply_chat_template(messages, return_tensors="pt").to(self.device)
+        # outputs = self.model.generate(encodeds, max_new_tokens=512, pad_token_id=self.tokenizer.eos_token_id)
+
+        # p=self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # print('---model output 2:\n',p)
+        transLTL=self.pattern.findall(p)[0]
+        if transLTL[-1]=='.':
+            transLTL=transLTL[:-1].strip()
+        else:
+            transLTL=transLTL.strip()
+        transLTL=self.ltlChecker.right_barkets_remover(transLTL)
+        print('transLTL:\n',transLTL)
+        return transLTL
+    def evaluate_model3(self, input_text):
+        # "LTL is a larger language model . . . . . . "
+        # self.pattern=re.compile("LTL is ([\S ]*)\.")
+        self.pattern=re.compile("LTL is ([^\.]*)\.")
+        messages=[
+        {"role": "user", "content": "translate natural description to linear temproal logic, first translate into a logical expression, and then translate into linear temproal logic, please pay specific attention to logic grammar, the natural language task is {}".format(input_text.strip())},
+        ]
+        encodeds = self.tokenizer.apply_chat_template(messages, return_tensors="pt").to(self.device)
+        outputs = self.model.generate(encodeds, max_new_tokens=512, pad_token_id=self.tokenizer.eos_token_id)
+        p=self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        print('---model output 1:\n',p)
+        # messages=[
+        #     {"role": "user", "content": "translate natural description to linear temproal logic, first translate into a logical expression, and then translate into linear temproal logic, the natural language task is {}".format(input_text.strip())},
+        #     {"role": "assistant", "content":p
+        #     },
+        #     {"role": "user", "content": " pay specific attention to brackets '()', given your linear temproal logic translation"},
+        # ]
+        
+        # encodeds = self.tokenizer.apply_chat_template(messages, return_tensors="pt").to(self.device)
+        # outputs = self.model.generate(encodeds, max_new_tokens=512, pad_token_id=self.tokenizer.eos_token_id)
+
+        # p=self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # print('---model output 2:\n',p)
+        transLTL=self.pattern.findall(p)
+        if len(transLTL)==0:
+            return False
+        transLTL=transLTL[0]
+        if transLTL[-1]=='.':
+            transLTL=transLTL[:-1].strip()
+        else:
+            transLTL=transLTL.strip()
+        transLTL=self.ltlChecker.right_barkets_remover(transLTL)
+        print('transLTL:\n',transLTL)
+        return transLTL
+    def translate(self,input_prompt:str=""):
+        print('input_prompt:\n',input_prompt)
+        replacer=self.replacer()
+        input_prompt=replacer.reTask2P(input_prompt)
+        # print(predicter( replace.reTask2P(input_prompt)))
+        # print(input_prompt)
+
+        
+        # print(p)
+        flag_check_false_count=0
+        flag_check=False
+        while not flag_check and flag_check_false_count<10:
+            flag_check_false_count+=1
+            flag_check=True
+            transLTL=self.evaluate_model3(input_prompt)
+            transLTL=transLTL.replace('Or','And')
+            transLTL=transLTL.replace('Globally','Finally')
+            if isinstance(transLTL,bool):
+                flag_check=False
+            elif not self.ltlChecker.AP_CorrCheck(input_prompt,transLTL):
+                print('AP_CorrCheck false')
+                flag_check=False
+            elif not self.ltlChecker.brackets_Check(transLTL):
+                print('brackets_Check false')
+                flag_check=False
+            # print(p)
+        return replacer.reP2Task(transLTL)
+
+
+if __name__=="__main__":
+    # translater=Mistral_NL2TL_translator()
+    # test_prompts=[
+    #     "Task_1.1.1 must precede Task_1.1.2, which in turn should precede Task_1.1.3, ",
+    #     "Task_1.1 must be completed before Task_1.2 starts, and Task_1.2 must be completed before Task_1.3 starts." ,
+    #     "Task_1.1 can be executed independently, after which Task_1.2 can be executed.",
+    #     "Task_1.2.4 must be completed first, followed by Task_1.2.2, then Task_1.2.3, and finally Task_1.2.1.",
+    #     "Task_1.2.4 is always executed first, followed by Task_1.2.3, then Task_1.2.2, and finally Task_1.2.1.",
+    #     "Task_1.2.1 and Task_1.2.2 can be executed independently, and both should eventually be completed.",
+    # ]
+    # for ret in test_prompts:
+    #     print(translater.translate(ret))    
+    #     print('\n','-'*20,'\n')
+    # exit()
+    class p2preplacer():
+        def reTask2P(self,input):
+            return input
+        def reP2Task(self,input):
+            return input
+    translater=Mistral_NL2TL_translator(replacer=p2preplacer)
+    import evaluate
+    import numpy as np
+    # from datasets import load_from_disk
+    from tqdm import tqdm
+
+    # Metric
+    metric = evaluate.load("rouge")
+    datapath='/home/user/xsj/NL2TL-dataset/collect2'
+    tokenized_dataset = load_dataset("json",  data_files={"train":os.path.join(datapath,"ltl_eng_train_mid_ascii_gptAuged.jsonl"),"test":os.path.join(datapath,"ltl_eng_test_mid_ascii_gptAuged.jsonl")})
+    print(tokenized_dataset)
+    # run predictions
+    # this can take ~45 minutes
+    import re 
+    # pattern=re.compile("\[Formal LTL\]:\n([\S ]*)\n")
+    predictions, references,input_sentence,output_sentence=[], [] , [], []
+    # with open()
+    for idx in range(len(tokenized_dataset['test']['natural'])):
+        # print(sample)
+        nl=tokenized_dataset['test']['natural'][idx]
+        transLTL=translater.translate(nl)
+        # p = translater.evaluate_model(nl)
+        # # print(p,l)
+        input_sentence.append(nl)
+
+        # transLTL=pattern.findall(p)
+        # # print(p)
+        predictions.append(transLTL) 
+        # output_sentence.append(p) 
+        # input_sentence.append(nl)
+        references.append(tokenized_dataset['test']['raw_ltl'][idx].strip())
+        print(idx,'\n',input_sentence[-1],
+            #   '\nout::\n',output_sentence[-1],
+              '\npre::\n',predictions[-1],
+              '\nref::\n',references[-1],'\n','-'*20,'\n')
+
+    # compute metric
+    rogue = metric.compute(predictions=predictions, references=references, use_stemmer=True)
+
+    # print results
+    print(f"Rogue1: {rogue['rouge1']* 100:2f}%")
+    print(f"rouge2: {rogue['rouge2']* 100:2f}%")
+    print(f"rougeL: {rogue['rougeL']* 100:2f}%")
+    print(f"rougeLsum: {rogue['rougeLsum']* 100:2f}%")
+    eval_output=np.array([input_sentence,predictions,references]).T
+    import pandas as pd 
+    eval_output=pd.DataFrame(eval_output)
+    pd.DataFrame.to_csv(eval_output,"/home/user/xsj/model_weight/mistral7b_mid_ascii_0327_eos_2aug1_quat8"+'/output')
+        # out llama
+        # Rogue1: 98.363321%
+        # rouge2: 95.987820%
+        # rougeL: 97.384820%
+        # rougeLsum: 97.382071%
+        
+        # this
+        # Rogue1: 98.543297%
+        # rouge2: 96.575248%
+        # rougeL: 97.720560%
+        # rougeLsum: 97.724880%
+    exit()
+    flag=True
+    while flag:
+        lines=[""]
+        try:
+            lines.append(input())
+            while True:
+                lines.append(input())
+        except:
+            pass
+        ret ="".join(lines)
+        print(ret)
+        if ret=="":
+            flag=False
+        
+        print(translater.translate(ret))
+
+

finetune/mistral7b/test.py

@@ -0,0 +1,32 @@
+import torch
+from transformers import (AutoModelForCausalLM,
+                         AutoTokenizer,
+                         BitsAndBytesConfig,
+                         TrainingArguments,
+                         pipeline,
+                         logging,
+                         TrainerCallback)
+device = "cuda" # the device to load the model onto
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit = True,
+    bnb_4bit_use_double_quant = False,
+    bnb_4bit_quant_type = 'nf4',
+    bnb_4bit_compute_dtype = getattr(torch, "float16")
+)
+model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2",quantization_config=bnb_config,)
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
+
+messages = [
+    {"role": "user", "content": "What is your favourite condiment?"},
+    {"role": "assistant", "content": "Well, I'm quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever I'm cooking up in the kitchen!"},
+    {"role": "user", "content": "Do you have mayonnaise recipes?"}
+]
+
+encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt")
+
+model_inputs = encodeds.to(device)
+# model.to(device)
+
+generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
+decoded = tokenizer.batch_decode(generated_ids)
+print(decoded[0])

finetune/realtime_run.py

@@ -0,0 +1,54 @@
+import torch
+from peft import PeftModel, PeftConfig
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+# Load peft config for pre-trained checkpoint etc.
+peft_model_id="finetuned_model/results"
+config = PeftConfig.from_pretrained(peft_model_id)
+
+# load base LLM model and tokenizer
+model = AutoModelForSeq2SeqLM.from_pretrained(config.base_model_name_or_path,  load_in_8bit=True,  device_map={"":0})
+tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+
+# Load the Lora model
+model = PeftModel.from_pretrained(model, peft_model_id, device_map={"":0})
+model.eval()
+
+print("Peft model loaded")
+
+from datasets import load_dataset
+from random import randrange
+
+
+import evaluate
+import numpy as np
+import datasets
+from tqdm import tqdm
+
+# Metric
+metric = evaluate.load("rouge")
+
+def evaluate_peft_model(sample,max_target_length=50):
+    # generate summary
+    outputs = model.generate(input_ids=sample["input_ids"].unsqueeze(0).cuda(), do_sample=True, top_p=0.9, max_new_tokens=max_target_length)
+    prediction = tokenizer.decode(outputs[0].detach().cpu().numpy(), skip_special_tokens=True)
+    # decode eval sample
+    # Replace -100 in the labels as we can't decode them.
+    labels = np.where(sample['labels'] != -100, sample['labels'], tokenizer.pad_token_id)
+    labels = tokenizer.decode(labels, skip_special_tokens=True)
+
+    # Some simple post-processing
+    return prediction, labels
+
+# load test dataset from distk
+# test_dataset = load_from_disk("data/eval/").with_format("torch")
+list_input = [{"natural": "go to P03 and then go to P04, remain in P04 until P05","raw_ltl":"0"}]
+test_dataset = datasets.Dataset.from_list(list_input)
+# run predictions
+# this can take ~45 minutes
+predictions, references = [] , []
+for sample in tqdm(test_dataset):
+    p,l = evaluate_peft_model(sample)
+    print(p,l)
+    predictions.append(p)
+    references.append(l)

finetune/test.py

@@ -0,0 +1,8 @@
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+tokenizer = AutoTokenizer.from_pretrained("bigscience/T0pp")
+model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp")
+
+inputs = tokenizer.encode("Is this review positive or negative? Review: this is the best cast iron skillet you will ever buy", return_tensors="pt")
+outputs = model.generate(inputs)
+print(tokenizer.decode(outputs[0]))