(model, dataset, output_path, device_num, prefix = None, suffix = None, chat = False)
| 771 | |
| 772 | |
| 773 | def infer_local_model(model, dataset, output_path, device_num, prefix = None, suffix = None, chat = False): |
| 774 | model = LocalModel(model, device_num, tensor_parallel_size = 8 if "Mixtral" in model else 4) |
| 775 | if dataset != "mbpp": |
| 776 | dataset = Dataset(dataset, data_path = os.path.join("test_datasets/", dataset.replace("/", "_")), full = False) |
| 777 | else: |
| 778 | dataset = Dataset(dataset, data_path = os.path.join("test_datasets/", dataset.replace("/", "_")), testfile_path = "datasets/MbppPlus.jsonl", full = False) |
| 779 | |
| 780 | results = {} |
| 781 | |
| 782 | finish = False |
| 783 | index = 0 |
| 784 | |
| 785 | while(not finish): |
| 786 | index += 1 |
| 787 | instance, finish = dataset.next() |
| 788 | if not chat: |
| 789 | prompt = dataset.get_prompt(instance) |
| 790 | if prefix != None or suffix != None: |
| 791 | prompt = (prefix if prefix != None else "") + prompt + (suffix if suffix != None else "") |
| 792 | else: |
| 793 | dialog = dataset.get_chat(instance) |
| 794 | try: |
| 795 | if not chat: |
| 796 | res = model.infer_one(prompt) |
| 797 | if isinstance(res, list): |
| 798 | results[prompt] = [res, True] |
| 799 | else: |
| 800 | results[prompt] = [res, False] |
| 801 | print(prompt) |
| 802 | print(res) |
| 803 | exit() |
| 804 | else: |
| 805 | res = model.infer_one_chat([dialog]) |
| 806 | if isinstance(res, list): |
| 807 | results[dialog[0]["content"]] = [res, True] |
| 808 | else: |
| 809 | results[dialog[0]["content"]] = [res, False] |
| 810 | print(dialog) |
| 811 | print(res) |
| 812 | exit() |
| 813 | |
| 814 | except Exception as e: |
| 815 | logger.error("Dataset: {}\nModel: {}\nPrompt:\n{}\n".format(dataset.name, model.name, prompt if not chat else dialog[0]["content"]) + str(e)) |
| 816 | results[prompt] = [str(e), False] |
| 817 | print("\r{}/{}".format(index, dataset.length()), end="", flush=True) |
| 818 | |
| 819 | res_dir_path = os.path.join(output_path, dataset.name.replace("/", "_")) |
| 820 | res_path = os.path.join(res_dir_path, model.name.replace("/", "_") + ".json") |
| 821 | |
| 822 | if not os.path.exists(res_dir_path): |
| 823 | os.mkdir(res_dir_path) |
| 824 | |
| 825 | with open(res_path, "w", encoding="utf-8") as of: |
| 826 | of.write(json.dumps(results, sort_keys=True, indent=4, separators=(',', ': '))) |
| 827 | |
| 828 | |
| 829 | def batch_infer_local_model(model, dataset, output_path, device_num, n = 200, temperature = 0.7, dataset_repo = "test_dataset", prefix = None, suffix = None, chat = False, context_length = None): |
nothing calls this directly
no test coverage detected