chore: import upstream snapshot with attribution

research/llm_embedder/scripts/llm-embedder.sh

@@ -0,0 +1,44 @@
+# the instruction and training config version
+version="llm-embedder"
+# the output folder
+output="llm-embedder"
+# the data root where you untar the data
+data_root="/data/llm-embedder"
+
+torchrun --nproc_per_node=8 run_dense.py --train_data \
+    llm-embedder:chat/msc/train.json \
+    llm-embedder:convsearch/qrecc/train.concat.json \
+    llm-embedder:lrlm/arxiv/train.json \
+    llm-embedder:lrlm/books3/train.json \
+    llm-embedder:lrlm/codeparrot/train.json \
+    llm-embedder:qa/msmarco/train.json \
+    llm-embedder:qa/nq/train.json \
+    llm-embedder:tool/toolbench/train.json \
+    llm-embedder:tool/toolbench/train.json \
+    llm-embedder:icl/icl/train.json \
+    --output_dir data/outputs/$output \
+    --save_steps 10000 \
+    --max_steps 10000 \
+    --logging_steps 100 \
+    --inbatch_same_dataset epoch \
+    --use_train_config \
+    --gradient_checkpointing \
+    --per_device_train_batch_size 100 \
+    --deepspeed data/deepspeed/stage0.json \
+    --version $version \
+    --learning_rate 5e-6 \
+    --data_root $data_root
+
+for model in "checkpoint-10000"
+do
+    torchrun --nproc_per_node 8 -m evaluation.eval_mmlu --query_encoder data/outputs/$output/$model/encoder --version $version --data_root $data_root
+    torchrun --nproc_per_node 8 -m evaluation.eval_popqa --query_encoder data/outputs/$output/$model/encoder --version $version --data_root $data_root
+    torchrun --nproc_per_node 8 -m evaluation.eval_msc --query_encoder data/outputs/$output/$model/encoder --version $version --data_root $data_root
+    torchrun --nproc_per_node 8 -m evaluation.eval_tool --query_encoder data/outputs/$output/$model/encoder --version $version --data_root $data_root
+    torchrun --nproc_per_node 8 -m evaluation.eval_lrlm --query_encoder data/outputs/$output/$model/encoder --eval_data llm-embedder:lrlm/books3/test.json --version $version --data_root $data_root
+    torchrun --nproc_per_node 8 -m evaluation.eval_lrlm --query_encoder data/outputs/$output/$model/encoder --eval_data llm-embedder:lrlm/arxiv/test.json --version $version --data_root $data_root
+    torchrun --nproc_per_node 8 -m evaluation.eval_lrlm --query_encoder data/outputs/$output/$model/encoder --eval_data llm-embedder:lrlm/codeparrot/test.json --version $version --data_root $data_root
+    torchrun --nproc_per_node 8 -m evaluation.eval_lrlm --query_encoder data/outputs/$output/$model/encoder --eval_data llm-embedder:lrlm/pg19/test.json --version $version --data_root $data_root
+    torchrun --nproc_per_node 8 -m evaluation.eval_icl --query_encoder data/outputs/$output/$model/encoder --version $version --data_root $data_root
+    torchrun --nproc_per_node 8 -m evaluation.eval_qrecc --query_encoder data/outputs/$output/$model/encoder --version $version --data_root $data_root
+done

research/llm_embedder/scripts/ours2st.py

@@ -0,0 +1,74 @@
+import os
+from typing import Optional, List
+from dataclasses import dataclass, field
+from sentence_transformers import models, SentenceTransformer
+from transformers import HfArgumentParser
+
+
+def convert_ours_ckpt_to_sentence_transformer(src_dir, dest_dir, pooling_method: List[str] = ['cls'], dense_metric: str="cos"):
+    assert os.path.exists(src_dir), f"Make sure the encoder path {src_dir} is valid on disk!"
+    assert "decoder" not in pooling_method, f"Pooling method 'decode' cannot be saved as sentence_transformers because it uses the decoder stack to produce sentence embedding."
+    if dest_dir is None:
+        dest_dir = src_dir
+
+    print(f"loading model from {src_dir} and saving the sentence_transformer model at {dest_dir}...")
+
+    word_embedding_model = models.Transformer(src_dir)
+    modules = [word_embedding_model]
+    ndim = word_embedding_model.get_word_embedding_dimension()
+
+    if "cls" in pooling_method:
+        pooling_model = models.Pooling(ndim, pooling_mode="cls")
+        pooling_method.remove("cls")
+    elif "mean" in pooling_method:
+        pooling_model = models.Pooling(ndim, pooling_mode="mean")
+        pooling_method.remove("mean")
+    else:
+        raise NotImplementedError(f"Fail to find cls or mean in pooling_method {pooling_method}!")
+    
+    modules.append(pooling_model)
+
+    if "dense" in pooling_method:
+        modules.append(models.Dense(ndim, ndim, bias=False))
+        pooling_method.remove("dense")
+    
+    assert len(pooling_method) == 0, f"Found unused pooling_method {pooling_method}!"
+
+    if dense_metric == "cos":
+        normalize_layer = models.Normalize()
+        modules.append(normalize_layer)
+
+    model = SentenceTransformer(modules=modules, device='cpu')
+    model.save(dest_dir)
+
+
+@dataclass
+class Args:
+    encoder: Optional[str] = field(
+        default=None,
+        metadata={'help': 'Path to the encoder model.'}
+    )
+    output_dir: Optional[str] = field(
+        default=None,
+        metadata={'help': 'Path to the output sentence_transformer model.'}
+    )
+    pooling_method: List[str] = field(
+        default_factory=lambda: ["cls"],
+        metadata={'help': 'Pooling methods to aggregate token embeddings for a sequence embedding. {cls, mean, dense, decoder}'}
+    )
+    dense_metric: str = field(
+        default="cos",
+        metadata={'help': 'What type of metric for dense retrieval? ip, l2, or cos.'}
+    )
+    model_cache_dir: Optional[str] = field(
+        default=None,
+        metadata={'help': 'Cache folder for huggingface transformers.'}
+    )
+
+    def __post_init__(self):
+        convert_ours_ckpt_to_sentence_transformer(self.encoder, self.output_dir, self.pooling_method, self.dense_metric)
+
+if __name__ == "__main__":
+    parser = HfArgumentParser([Args])
+    args, = parser.parse_args_into_dataclasses()
+