update readme

Author: 545999961

research/LM_Cocktail/README.md

@@ -49,7 +49,7 @@ The merged model can be used to perform multiple tasks.
 Install the latest version from source (Recommended): 
 ```bash
 git clone https://github.com/FlagOpen/FlagEmbedding.git
-cd FlagEmbedding/LM_Cocktail
+cd FlagEmbedding/research/LM_Cocktail
 pip install -e .
 ```
 Install by pip:
@@ -260,6 +260,7 @@ torchrun --nproc_per_node 8 -m evaluation.eval_mmlu \
 - Models: we fine-tune the [bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5) on 9 tasks, and you can find the fine-tuned models at this [link](https://huggingface.co/Shitao).
 - Examples Data: [./embedder_examples.json]()
 
+
 Use [MTEB script](https://github.com/FlagOpen/FlagEmbedding/tree/master/C_MTEB) to evaluate the mixed embedding model:
 ```bash
 python eval_MTEB.py --model_name_or_path mixed_model --task_type Retrieval

research/llm_embedder/README.md

@@ -20,19 +20,51 @@ This is the codebase for LLM-Embedder, a unified embedding model to comprehensiv
 ### Using `FlagEmbedding`
 ```pip install -U FlagEmbedding```
 ```python
-from FlagEmbedding import LLMEmbedder
+from FlagEmbedding import FlagModel
+
+INSTRUCTIONS = {
+    "qa": {
+        "query": "Represent this query for retrieving relevant documents: ",
+        "key": "Represent this document for retrieval: ",
+    },
+    "icl": {
+        "query": "Convert this example into vector to look for useful examples: ",
+        "key": "Convert this example into vector for retrieval: ",
+    },
+    "chat": {
+        "query": "Embed this dialogue to find useful historical dialogues: ",
+        "key": "Embed this historical dialogue for retrieval: ",
+    },
+    "lrlm": {
+        "query": "Embed this text chunk for finding useful historical chunks: ",
+        "key": "Embed this historical text chunk for retrieval: ",
+    },
+    "tool": {
+        "query": "Transform this user request for fetching helpful tool descriptions: ",
+        "key": "Transform this tool description for retrieval: "
+    },
+    "convsearch": {
+        "query": "Encode this query and context for searching relevant passages: ",
+        "key": "Encode this passage for retrieval: ",
+    },
+}
 
 # Define queries and keys
 queries = ["test query 1", "test query 2"]
 keys = ["test key 1", "test key 2"]
 
-# Load model (automatically use GPUs)
-model = LLMEmbedder('BAAI/llm-embedder', use_fp16=False)
-
 # Encode for a specific task (qa, icl, chat, lrlm, tool, convsearch)
 task = "qa"
-query_embeddings = model.encode_queries(queries, task=task)
-key_embeddings = model.encode_keys(keys, task=task)
+
+# Load model (automatically use GPUs)
+model = FlagModel('BAAI/llm-embedder', 
+                  use_fp16=False,
+                  query_instruction_for_retrieval=INSTRUCTIONS[task]['query'],
+                  passage_instruction_for_retrieval=INSTRUCTIONS[task]['key'],
+                  devices=['cuda:0'])
+
+query_embeddings = model.encode_queries(queries)
+key_embeddings = model.encode_corpus(keys)
 
 similarity = query_embeddings @ key_embeddings.T
 print(similarity)

research/old-examples/pretrain/README.md

@@ -11,15 +11,9 @@ pip install -U FlagEmbedding
 * **from source**
 ```
 git clone https://github.com/FlagOpen/FlagEmbedding.git
-cd FlagEmbedding
-pip install  .
-```
-For development, install as editable:
-```
-pip install -e .
+cd FlagEmbedding/research/old-examples/pretrain
 ```
 
-
 ## 2. Data format
 Train data should be a json file, where each line is a dict like this:
 ```
@@ -31,7 +25,7 @@ See [toy_pretrain_data.jsonl](https://github.com/FlagOpen/FlagEmbedding/blob/mas
 
 ```bash
 torchrun --nproc_per_node {number of gpus} \
--m FlagEmbedding.baai_general_embedding.retromae_pretrain.run \
+-m retromae_pretrain.run \
 --output_dir {path to save model} \
 --model_name_or_path BAAI/bge-large-en \
 --train_data toy_pretrain_data.jsonl \
@@ -47,4 +41,3 @@ torchrun --nproc_per_node {number of gpus} \
 More training arguments please refer to [transformers.TrainingArguments](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments). 
 After training, the encoder model will saved to `{output_dir}/encoder_model`
 
-

research/old-examples/pretrain/retromae_pretrain/__init__.py

@@ -0,0 +1,2 @@
+
+

research/old-examples/pretrain/retromae_pretrain/arguments.py

@@ -0,0 +1,43 @@
+import os
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+@dataclass
+class DataTrainingArguments:
+    train_data: Optional[str] = field(
+        default=None, metadata={"help": "Path to pretrain data"}
+    )
+    tokenizer_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+    )
+    max_seq_length: Optional[int] = field(
+        default=512,
+        metadata={
+            "help": "The maximum total input sequence length after tokenization. Sequences longer "
+                    "than this will be truncated. Default to the max input length of the model."
+        },
+    )
+    encoder_mlm_probability: float = field(default=0.3, metadata={"help": "mask ratio for encoder"})
+    decoder_mlm_probability: float = field(default=0.5, metadata={"help": "mask ratio for decoder"})
+
+    def __post_init__(self):
+        if not os.path.exists(self.train_data):
+            raise FileNotFoundError(f"cannot find file: {self.train_data}, please set a true path")
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
+    """
+    model_name_or_path: Optional[str] = field(
+        default='bert-base-uncased',
+        metadata={
+            "help": "The model checkpoint for weights initialization."
+                    "Don't set if you want to train a model from scratch."
+        },
+    )
+    config_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+    )

research/old-examples/pretrain/retromae_pretrain/data.py

@@ -0,0 +1,100 @@
+import os
+import random
+from copy import deepcopy
+from dataclasses import dataclass
+
+import torch.utils.data.dataset
+from datasets import Dataset, load_dataset, concatenate_datasets
+from transformers import DataCollatorForWholeWordMask
+
+from .utils import tensorize_batch
+
+
+class DatasetForPretraining(torch.utils.data.Dataset):
+    def __init__(self, data_dir):
+        if os.path.isdir(data_dir):
+            datasets = []
+            for file in os.listdir(data_dir):
+                print(f"Loading {file}")
+                file = os.path.join(data_dir, file)
+                datasets.append(self.load_dataset(file))
+            self.dataset = concatenate_datasets(datasets)
+        else:
+            print(f"Loading {data_dir}")
+            self.dataset = self.load_dataset(data_dir)
+
+    def load_dataset(self, file):
+        if file.endswith('.jsonl') or file.endswith('.json'):
+            return load_dataset('json', data_files=file)['train']
+        elif os.path.isdir(file):
+            return Dataset.load_from_disk(file)
+        else:
+            raise NotImplementedError(f"Not support this file format:{file}")
+
+    def __getitem__(self, item):
+        return self.dataset[item]['text']
+
+    def __len__(self):
+        return len(self.dataset)
+
+
+@dataclass
+class RetroMAECollator(DataCollatorForWholeWordMask):
+    max_seq_length: int = 512
+    encoder_mlm_probability: float = 0.15
+    decoder_mlm_probability: float = 0.15
+
+    def __call__(self, examples):
+        input_ids_batch = []
+        attention_mask_batch = []
+        encoder_mlm_mask_batch = []
+        decoder_labels_batch = []
+        decoder_matrix_attention_mask_batch = []
+
+        for e in examples:
+
+            e_trunc = self.tokenizer.encode(e, max_length=self.max_seq_length, truncation=True)
+            tokens = [self.tokenizer._convert_id_to_token(tid) for tid in e_trunc]
+
+            self.mlm_probability = self.encoder_mlm_probability
+            text_encoder_mlm_mask = self._whole_word_mask(tokens)
+
+            self.mlm_probability = self.decoder_mlm_probability
+            mask_set = []
+            for _ in range(min(len(tokens), 128)):
+                mask_set.append(self._whole_word_mask(tokens))
+
+            text_matrix_attention_mask = []
+            for i in range(len(tokens)):
+                idx = random.randint(0, min(len(tokens), 128) - 1)
+                text_decoder_mlm_mask = deepcopy(mask_set[idx])
+                text_decoder_mlm_mask[i] = 1
+                text_matrix_attention_mask.append(text_decoder_mlm_mask)
+
+            input_ids_batch.append(torch.tensor(e_trunc))
+            attention_mask_batch.append(torch.tensor([1] * len(e_trunc)))
+            e_trunc[0] = -100
+            e_trunc[-1] = -100
+            decoder_labels_batch.append(torch.tensor(e_trunc))
+
+            encoder_mlm_mask_batch.append(torch.tensor(text_encoder_mlm_mask))
+            decoder_matrix_attention_mask_batch.append(1 - torch.tensor(text_matrix_attention_mask))
+
+        input_ids_batch = tensorize_batch(input_ids_batch, self.tokenizer.pad_token_id)
+        attention_mask_batch = tensorize_batch(attention_mask_batch, 0)
+        origin_input_ids_batch = input_ids_batch.clone()
+        encoder_mlm_mask_batch = tensorize_batch(encoder_mlm_mask_batch, 0)
+        encoder_input_ids_batch, encoder_labels_batch = self.torch_mask_tokens(input_ids_batch, encoder_mlm_mask_batch)
+        decoder_labels_batch = tensorize_batch(decoder_labels_batch, -100)
+        matrix_attention_mask_batch = tensorize_batch(decoder_matrix_attention_mask_batch, 0)
+
+        batch = {
+            "encoder_input_ids": encoder_input_ids_batch,
+            "encoder_attention_mask": attention_mask_batch,
+            "encoder_labels": encoder_labels_batch,
+            "decoder_input_ids": origin_input_ids_batch,
+            "decoder_attention_mask": matrix_attention_mask_batch,  # [B,L,L]
+            "decoder_labels": decoder_labels_batch,
+        }
+
+        return batch