first commit

update readme

update readme

update readme

Update benchmarks.py

Update download_adapters.sh

Update esft.py

scripts/download_adapters.sh

@@ -0,0 +1,22 @@
+tasks=(
+    math
+    code
+    intent
+    summary
+    law
+    translation
+)
+
+mkdir -p all_models/adapters/token
+git lfs install
+for i in {0..5}
+do
+    git clone https://huggingface.co/deepseek-ai/ESFT-token-${tasks[$i]}-lite ./all_models/adapters/token/${tasks[$i]}
+done
+
+
+mkdir -p all_models/adapters/gate
+for i in {0..5}
+do
+    git clone https://huggingface.co/deepseek-ai/ESFT-gate-${tasks[$i]}-lite ./all_models/adapters/gate/${tasks[$i]}
+done

scripts/eval.py

@@ -0,0 +1,53 @@
+import json
+import argparse
+from benchmarks import *
+import os
+from esft import load_base_model, add_adapter
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--adapter_dir", type=str, required=True)
+parser.add_argument("--base_model_path", type=str, required=True)
+parser.add_argument("--eval_datasets", type=str, required=True)
+parser.add_argument("--output_dir", type=str, required=True)
+parser.add_argument("--max_new_tokens", type=int, default=128)
+parser.add_argument("--eval_batch_size", type=int, default=16)
+parser.add_argument("--openai_api_key", type=str, required=True)
+parser.add_argument("--debug", action='store_true')
+args = parser.parse_args()
+
+base_model_path = args.base_model_path
+adapter_dir = args.adapter_dir
+eval_datasets = args.eval_datasets.split(",")
+
+config = {"max_new_tokens": args.max_new_tokens, "eval_batch_size": args.eval_batch_size, "openai_api_key": args.openai_api_key}
+
+evaluator_map={"intent": IntentEvaluator, "summary": SummaryEvaluator, "law": LawEvaluator, "translation": TranslationEvaluator}
+
+print("Loading base model...")
+model, tokenizer = load_base_model(base_model_path)
+
+for dataset_name in eval_datasets:
+    print(f"Running evaluation on {dataset_name}...")
+    dataset = [json.loads(i) for i in open(f"datasets/eval/{dataset_name}.jsonl").readlines()]
+    if args.debug:
+        print("Debugging. Shortening the dataset length")
+        dataset = dataset[:16]
+
+    evaluator = evaluator_map[dataset_name](dataset, config)
+    print("Adding adapter...")
+    model.model, original_state_dict = add_adapter(model.model, os.path.join(adapter_dir, dataset_name), return_original_states=True) # add adapter to model and convert original states to buffer.
+    results, metrics = evaluator.evaluate(model, tokenizer)
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    with open(os.path.join(args.output_dir, dataset_name + ".jsonl"), "w") as f:
+        for res, m in zip(results, metrics):
+            obj = {
+                "example": res,
+                "score": m
+            }
+            f.write(json.dumps(obj, ensure_ascii=False) + "\n")
+
+    model.model.load_state_dict(original_state_dict) # convert to original model
+
+
+

scripts/eval.sh

@@ -0,0 +1,12 @@
+# first, download adapter models and put them to the corresponding directories
+
+
+python scripts/eval.py \
+    --eval_datasets=translation \
+    --base_model_path=deepseek-ai/ESFT-vanilla-lite \
+    --adapter_dir=all_models/adapters/token \
+    --output_dir=results/completions/token \
+    --max_new_tokens=512 \
+    --openai_api_key=REPLACE_WITH_YOUR_KEY \
+    --eval_batch_size=2
+

scripts/generate_expert_config.py

@@ -0,0 +1,50 @@
+import argparse
+import json
+import os
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--eval_datasets", type=str, required=True)
+parser.add_argument("--expert_scores_dir", type=str, required=True)
+parser.add_argument("--output_dir", type=str, required=True)
+parser.add_argument("--score_function", type=str, required=True)
+parser.add_argument("--top_p", type=float, required=True)
+parser.add_argument("--train_shared_experts", action="store_true")
+parser.add_argument("--train_non_expert_modules", action="store_true")
+
+args = parser.parse_args()
+
+eval_datasets = args.eval_datasets.split(",")
+expert_scores_dir = args.expert_scores_dir
+output_dir = args.output_dir
+score_function = args.score_function
+top_p = args.top_p
+train_shared_experts = args.train_shared_experts
+train_non_expert_modules = args.train_non_expert_modules
+
+for dataset_name in eval_datasets:
+    summary_file = f"{expert_scores_dir}/{dataset_name}/summary.json"
+    expert_cfg = {"experts": {}, "shared_experts": train_shared_experts, "non_expert_modules": train_non_expert_modules}
+
+    with open(summary_file) as f:
+        data = json.load(f)
+        assert score_function in ["gate", "token"], f"Unknown score function: {score_function}"
+        scores = data[f"{score_function}_scores"]
+
+        for layer, l_score in scores.items():
+            l_score = [(int(k), v) for k,v in l_score.items()]
+            l_score = sorted(l_score, key=lambda x: x[1], reverse=True)
+            # get the top experts that make the threshold exceed top_p
+            selected_experts = []
+            current_score = 0
+            for expert, score in l_score:
+                if current_score >= top_p:
+                    break
+                selected_experts.append(expert)
+                current_score += score
+            expert_cfg["experts"][layer] = selected_experts
+    
+    os.makedirs(output_dir, exist_ok=True)
+    with open(f"{output_dir}/{dataset_name}.json", "w") as f:
+        json.dump(expert_cfg, f)
+
+        

scripts/generate_expert_config.sh

@@ -0,0 +1,14 @@
+python scripts/get_expert_scores.py \
+    --eval_datasets=intent,summary,law,translation \
+    --base_model_path=deepseek-ai/ESFT-vanilla-lite \
+    --output_dir=results/expert_scores \
+    --n_sample_tokens=8192 # this sample size is a hyperparameter
+
+python scripts/generate_expert_config.py \
+    --eval_datasets=intent,summary,law,translation \
+    --expert_scores_dir=results/expert_scores \
+    --output_dir=results/expert_configs \
+    --score_function=token \
+    --top_p=0.2 # the scoring function and top_p are hyperparameters
+    # --train_shared_experts
+    # --train_non_expert_modules

scripts/get_expert_scores.py

@@ -0,0 +1,75 @@
+import json
+from benchmarks import *
+import os
+import torch
+from torch import nn
+import argparse
+from random import shuffle
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from utils import get_formatted_input_and_target
+
+# constants for deepseek-v2-lite
+TOP_K=6
+N_EXPERTS=64
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--base_model_path", type=str, required=True)
+parser.add_argument("--eval_datasets", type=str, required=True)
+parser.add_argument("--output_dir", type=str, required=True)
+parser.add_argument("--n_sample_tokens", type=int, required=True)
+args = parser.parse_args()
+
+eval_datasets = args.eval_datasets.split(",")
+output_dir = args.output_dir
+base_model_path = args.base_model_path
+n_sample_tokens = args.n_sample_tokens
+
+model, tokenizer = AutoModelForCausalLM.from_pretrained(base_model_path, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto"), AutoTokenizer.from_pretrained(base_model_path)
+model.config.log_expert_weights = True
+
+for dataset_name in eval_datasets:
+    dataset = [json.loads(i) for i in open(f"datasets/train/{dataset_name}.jsonl").readlines()]
+    shuffle(dataset)
+    model.config.expert_log_dir = os.path.join(args.output_dir, dataset_name)
+    # make dir -p this
+    os.makedirs(os.path.join(args.output_dir, dataset_name), exist_ok=True)
+    done_tokens = 0
+    for instance in dataset:
+        input_ids, target_ids = get_formatted_input_and_target(instance['messages'], tokenizer, -100)
+        model(input_ids=torch.tensor(input_ids).unsqueeze(0), labels=torch.tensor(target_ids).unsqueeze(0))
+        done_tokens += len(input_ids)
+        if done_tokens >= n_sample_tokens:
+            break
+
+    # open all files under os.path.join(args.output_dir, dataset_name). For each file, generate a summary of it
+    # and write it to a file in the same directory
+    files = os.listdir(os.path.join(args.output_dir, dataset_name))
+    summary_file = os.path.join(args.output_dir, dataset_name, "summary.json")
+    token_scores = {}
+    gate_scores = {}
+
+    for file in files:
+        if not file.endswith(".txt"):
+            continue
+        layer_idx = file.split("_")[2].split(".")[0]
+        token_scores[layer_idx] = {expert:0 for expert in range(N_EXPERTS)}
+        gate_scores[layer_idx] = {expert:0 for expert in range(N_EXPERTS)}
+
+        with open(os.path.join(args.output_dir, dataset_name, file)) as f:
+            data = f.readlines()
+            for line in data:
+                expert_ids, expert_weights = line.split("\t\t")
+                expert_ids = [int(i) for i in expert_ids.split("\t")]
+                expert_weights = [float(i) for i in expert_weights.split("\t")]
+                for expert_id, expert_weight in zip(expert_ids, expert_weights):
+                    gate_scores[layer_idx][expert_id] += expert_weight
+                    token_scores[layer_idx][expert_id] += 1. / TOP_K
+            total = sum(token_scores[layer_idx].values())
+            gate_scores[layer_idx] = {expert: round(gate_scores[layer_idx][expert] / total, 4) for expert in gate_scores[layer_idx]}
+            token_scores[layer_idx] = {expert: round(token_scores[layer_idx][expert] / total, 4) for expert in token_scores[layer_idx]}
+
+
+    with open(summary_file, "w") as f:
+        f.write(json.dumps({"token_scores": token_scores, "gate_scores": gate_scores}))
+
+