File size: 1,177 Bytes
f021684 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
from LM_Cocktail import mix_models_by_layers
import argparse
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model_type", type=str, default="decoder", help="Type of model to be mixed")
parser.add_argument("--output_path", type=str, default="./mixed_llm", help="Path to save the mixed model")
parser.add_argument("--max_length", type=int, default=100, help="Maximum length of the sequence to be generated")
parser.add_argument("--models", type=str, nargs='+', default=["meta-llama/Llama-2-7b-chat-hf", "Shitao/llama2-ag-news"], help="Path to the models to be mixed")
parser.add_argument("--weights", type=float, nargs='+', default=[0.7, 0.3], help="Weights of the models to be mixed")
parser.add_argument("--save_precision", type=str, default='float32', help="mixed model saved format")
args = parser.parse_args()
# Mix Large Language Models (LLMs) and save the combined model to the path: ./mixed_llm
model = mix_models_by_layers(
model_names_or_paths=args.models,
model_type=args.model_type,
weights=args.weights,
output_path=args.output_path)
print(model) |