36 | #
37 | # ```
38 | # This function convert it into the correct markdown format
39 | return re.sub(code_lang_pattern, code_lang_format, val)
40 |
41 |
42 | def html_to_markdown(val: str) -> str:
43 | # Remove all . This is required to make intent work in code blocks.
44 | val = re.sub(div_pattern, "", val)
45 | # Remove all . This is required to make underscores work in code blocks.
46 | val = re.sub(span_pattern, "", val)
47 | # Markdown to html
48 | val = markdownify.markdownify(val).strip()
49 | # Reformat code
50 | val = reformat_code(val)
51 |
52 | # Remove noisy "[number] / [number]" at the beginning
53 | noise = re.search(regenerate_pattern, val)
54 | if noise and noise.start() == 0:
55 | val = val[noise.end() :]
56 | # Remove noisy "Copy[number] chars / [number] words"
57 | val = re.sub(copy_chars_pattern, "", val)
58 | # Remove empty code block ```\nCopy code\n```
59 | val = re.sub(copy_code_pattern, "", val)
60 |
61 | # Strip
62 | val = val.replace("\n\n\n", "\n").strip()
63 |
64 | return val
65 |
66 |
67 | def contain_blocked_words(val: str) -> bool:
68 | blocked_words = ["openai", "chatgpt"]
69 | for w in blocked_words:
70 | if w in val.lower():
71 | return True
72 | return False
73 |
74 |
75 | def clean_html_one_sample(sample):
76 | roles = ["human", "gpt"]
77 |
78 | if len(sample["conversations"]) <= 1:
79 | return (sample, 1)
80 |
81 | # Adjust the offset for cases like https://sharegpt.com/c/VyaZlh4
82 | if sample["conversations"][0]["from"] != "human":
83 | sample["conversations"] = sample["conversations"][1:]
84 | if len(sample["conversations"]) <= 1:
85 | return (sample, 1)
86 |
87 | if sample["conversations"][-1]["from"] == "human":
88 | sample["conversations"] = sample["conversations"][:-1]
89 | if len(sample["conversations"]) <= 1:
90 | return (sample, 1)
91 |
92 | for i, c in enumerate(sample["conversations"]):
93 | if c["from"] != roles[i % 2]:
94 | return (sample, 2)
95 |
96 | if contain_blocked_words(c["value"]):
97 | return (sample, 3)
98 |
99 | try:
100 | new_val = html_to_markdown(c["value"])
101 | except (bs4.builder.ParserRejectedMarkup, AssertionError):
102 | return (sample, 4)
103 | new_val = converter.convert(new_val)
104 | c["value"] = new_val
105 |
106 | return (sample, 0)
107 |
108 |
109 | def clean_html_all(content, begin, end):
110 | """
111 | Clean the source html files.
112 | """
113 | cnt_skip = 0
114 | cnt_blocked_words = 0
115 | cnt_wrong_format = 0
116 | cnt_parser_error = 0
117 | cnt_too_short = 0
118 | cnt_id_duplication = 0
119 | cnt_value_duplication = 0
120 | cnt_tag = 0
121 |
122 | content = content[begin:end]
123 | processed = []
124 | with ProcessPoolExecutor() as executor:
125 | for result in tqdm(
126 | executor.map(clean_html_one_sample, content), total=len(content)
127 | ):
128 | processed.append(result)
129 |
130 | visited = {}
131 | new_content = []
132 | for sample, error_code in tqdm(processed):
133 | cid = sample["id"]
134 | skipped = True
135 |
136 | if error_code != 0:
137 | if error_code == 1:
138 | print(f"id {cid} is too short")
139 | cnt_too_short += 1
140 | elif error_code == 2:
141 | print(f"id {cid} has a wrong format")
142 | cnt_wrong_format += 1
143 | elif error_code == 3:
144 | print(f"id {cid} contains blocked words")
145 | cnt_blocked_words += 1
146 | elif error_code == 4:
147 | print(f"id {cid} contains parser errors")
148 | cnt_parser_error += 1
149 | else:
150 | raise ValueError(f"Invalid error_code: {error_code}")
151 | elif cid in visited:
152 | print(f"id {cid} is an id duplication of {visited[cid]}")
153 | cnt_id_duplication += 1
154 | elif (
155 | sample["conversations"][1]["value"],
156 | len(sample["conversations"]),
157 | ) in visited:
158 | key = (sample["conversations"][1]["value"], len(sample["conversations"]))
159 | print(f"id {cid} is a value duplication of {visited[key]}")
160 | cnt_value_duplication += 1
161 | else:
162 | key = (sample["conversations"][1]["value"], len(sample["conversations"]))
163 | visited[cid] = visited[key] = cid
164 | skipped = False
165 |
166 | if not skipped:
167 | new_content.append(sample)
168 | else:
169 | cnt_skip += 1
170 |
171 | print(
172 | f"total: {len(content)}, skip: {cnt_skip}, new: {len(new_content)}, "
173 | f"cnt_blocked_words: {cnt_blocked_words}, cnt_parser_error: {cnt_parser_error}, "
174 | f"cnt_wrong_format: {cnt_wrong_format}, "
175 | f"cnt_too_short: {cnt_too_short}, cnt_id_duplication: {cnt_id_duplication}, "
176 | f"cnt_value_duplication: {cnt_value_duplication}, "
177 | )
178 |
179 | return new_content
180 |
181 |
182 | def main(args):
183 | content = json.load(open(args["in_file"], "r",encoding='utf-8'))
184 | content = clean_html_all(content, args["begin"], args["end"])
185 | json_content = json.dumps(content, ensure_ascii=False,indent=2)
186 | with open(args["out_file"], 'w', encoding='utf-8') as f:
187 | f.write(json_content)
188 |
189 |
190 | if __name__ == "__main__":
191 | parser = argparse.ArgumentParser()
192 | parser.add_argument("--in-file", type=str, required=True)
193 | parser.add_argument("--out-file", type=str, default="sharegpt_clean.json")
194 | parser.add_argument("--begin", type=int)
195 | parser.add_argument("--end", type=int)
196 | parser.add_argument("--debug", action="store_true")
197 | args = parser.parse_args()
198 | main(vars(args))
199 |
--------------------------------------------------------------------------------
/fastchat/data/merge.py:
--------------------------------------------------------------------------------
1 | """
2 | Merge two conversation files into one
3 |
4 | Usage: python3 -m fastchat.data.merge --in file1.json file2.json --out merged.json
5 | """
6 |
7 | import argparse
8 | import json
9 | from typing import Dict, Sequence, Optional
10 |
11 |
12 | if __name__ == "__main__":
13 | parser = argparse.ArgumentParser()
14 | parser.add_argument("--in-file", type=str, required=True, nargs="+")
15 | parser.add_argument("--out-file", type=str, default="merged.json")
16 | args = parser.parse_args()
17 |
18 | new_content = []
19 | for in_file in args.in_file:
20 | content = json.load(open(in_file, "r",encoding='utf-8'))
21 | new_content.extend(content)
22 |
23 | json.dump(new_content, open(args.out_file, "w",encoding='utf-8'), ensure_ascii=False,indent=2)
24 |
--------------------------------------------------------------------------------
/fastchat/data/optional_clean.py:
--------------------------------------------------------------------------------
1 | """
2 | Do optional cleaning (e.g., remove some languages).
3 |
4 | Usage:
5 | python3 -m fastchat.data.optional_clean --in input.json --out output.json --keep-lang en
6 | python3 -m fastchat.data.optional_clean --in input.json --out output.json --skip-lang en
7 |
8 | Requirement:
9 | pip3 install polyglot icu pyicu pycld2 morfessor
10 | """
11 | import argparse
12 | import json
13 | import re
14 |
15 | import polyglot
16 | from polyglot.detect import Detector
17 | import pycld2
18 | from tqdm import tqdm
19 |
20 |
21 | def skip(conv, args):
22 | # Remove certain languages
23 | if args.keep_lang != "all" or args.skip_lang is not None:
24 | text = "\n".join([x["value"] for x in conv["conversations"]])
25 | try:
26 | lang_code = Detector(text).language.code
27 | except (pycld2.error, polyglot.detect.base.UnknownLanguage):
28 | lang_code = "unknown"
29 |
30 | if args.keep_lang != "all" and lang_code != args.keep_lang:
31 | return True
32 |
33 | if lang_code == args.skip_lang:
34 | return True
35 |
36 | # Remove repetitive numbers
37 | if args.reduce_rep:
38 | for sentence in conv["conversations"]:
39 | val = sentence["value"]
40 | sub = re.search(r"(\d)\1{8}", val)
41 | if sub is not None:
42 | return True
43 |
44 | return False
45 |
46 |
47 | if __name__ == "__main__":
48 | parser = argparse.ArgumentParser()
49 | parser.add_argument("--in-file", type=str, required=True)
50 | parser.add_argument("--out-file", type=str)
51 | parser.add_argument(
52 | "--keep-lang",
53 | type=str,
54 | default="all",
55 | choices=["all", "en"],
56 | help="Only keep certain langauges.",
57 | )
58 | parser.add_argument("--skip-lang", type=str, help="Skip a specific language.")
59 | # NOTE: Be careful about reduce_rep which may remove some good data.
60 | # For example, addresses could have long consecutive 0's
61 | parser.add_argument("--reduce-rep", action="store_true")
62 | args = parser.parse_args()
63 |
64 | in_file = args.in_file
65 | out_file = args.out_file
66 | keep_lang = args.keep_lang
67 | skip_lang = args.skip_lang
68 | reduce_rep = args.reduce_rep
69 | assert keep_lang == "all" or skip_lang is None
70 |
71 | if out_file is None:
72 | out_file = "sharegpt_clean"
73 | if keep_lang != "all":
74 | out_file += "_" + keep_lang
75 | if skip_lang is not None:
76 | out_file += "_skip_" + skip_lang
77 | if reduce_rep:
78 | out_file += "_reduce_rep"
79 | out_file += ".json"
80 |
81 | content = json.load(open(in_file, "r",encoding='utf-8'))
82 | num_conv = len(content)
83 |
84 | new_content = []
85 | for conv in tqdm(content):
86 | if not skip(conv, args):
87 | new_content.append(conv)
88 |
89 | print(f"return {len(new_content)} out of {len(content)}, start dump ...")
90 | json_content = json.dumps(new_content, ensure_ascii=False,indent=2)
91 | with open(out_file, 'w', encoding='utf-8') as f:
92 | f.write(json_content)
93 |
--------------------------------------------------------------------------------
/fastchat/data/split_long_conversation.py:
--------------------------------------------------------------------------------
1 | """
2 | Split long conversations based on certain max length.
3 |
4 | Usage: python3 -m fastchat.data.split_long_conversation \
5 | --in sharegpt_clean.json \
6 | --out sharegpt_split.json \
7 | --model-name-or-path $
8 | """
9 | import argparse
10 | from concurrent.futures import ProcessPoolExecutor
11 | import json
12 | from typing import Dict, Sequence, Optional
13 |
14 | import transformers
15 | from tqdm import tqdm
16 |
17 |
18 | def make_sample(sample, start_idx, end_idx):
19 | assert (end_idx - start_idx) % 2 == 0
20 | return {
21 | "id": sample["id"] + "_" + str(start_idx),
22 | "conversations": sample["conversations"][start_idx:end_idx],
23 | }
24 |
25 |
26 | tokenizer = max_length = None
27 |
28 |
29 | def split_one_sample(sample):
30 | tokenized_lens = []
31 | conversations = sample["conversations"]
32 | conversations = conversations[: len(conversations) // 2 * 2]
33 | for c in conversations:
34 | length = len(tokenizer(c["value"]).input_ids) + 6
35 | tokenized_lens.append(length)
36 |
37 | start_idx = 0
38 | cur_len = 0
39 |
40 | if len(conversations) % 2 != 0 or len(conversations) < 2:
41 | return []
42 |
43 | new_samples = []
44 | for i in range(0, len(conversations), 2):
45 | tmp_len = tokenized_lens[i] + tokenized_lens[i + 1]
46 | if cur_len + tmp_len > max_length:
47 | new_samples.append(make_sample(sample, start_idx, i))
48 | start_idx = i
49 | cur_len = 0
50 | elif i == len(conversations) - 2:
51 | new_samples.append(make_sample(sample, start_idx, i + 2))
52 |
53 | cur_len += tmp_len
54 |
55 | return new_samples
56 |
57 |
58 | def split_all(content, begin, end, tokenizer_, max_length_):
59 | """
60 | Keep the maximum round of conversations within the max token length constraint
61 | """
62 | global tokenizer, max_length
63 | tokenizer = tokenizer_
64 | max_length = max_length_
65 |
66 | content = content[begin:end]
67 | new_content = []
68 |
69 | with ProcessPoolExecutor() as executor:
70 | for result in tqdm(executor.map(split_one_sample, content), total=len(content)):
71 | new_content.extend(result)
72 |
73 | return new_content
74 |
75 |
76 | def filter_invalid_roles(content):
77 | new_content = []
78 | for i, c in enumerate(content):
79 | roles = ["human", "gpt"]
80 | if len(c["conversations"]) <= 0:
81 | continue
82 |
83 | valid = True
84 | for j, s in enumerate(c["conversations"]):
85 | if s["from"] != roles[j % 2]:
86 | valid = False
87 | break
88 |
89 | if valid:
90 | new_content.append(c)
91 |
92 | return new_content
93 |
94 |
95 | def main(args):
96 | content = json.load(open(args.in_file, "r"))
97 | tokenizer = transformers.AutoTokenizer.from_pretrained(
98 | args.model_name_or_path,
99 | model_max_length=args.max_length,
100 | padding_side="right",
101 | use_fast=False,
102 | )
103 | new_content = split_all(content, args.begin, args.end, tokenizer, args.max_length)
104 | new_content = filter_invalid_roles(new_content)
105 |
106 | print(f"total: {len(content)}, new: {len(new_content)}")
107 | json_content = json.dumps(new_content, ensure_ascii=False,indent=2)
108 | with open(args.out_file, 'w', encoding='utf-8') as f:
109 | f.write(json_content)
110 |
111 |
112 | if __name__ == "__main__":
113 | parser = argparse.ArgumentParser()
114 | parser.add_argument("--in-file", type=str, required=True)
115 | parser.add_argument("--out-file", type=str, default="sharegpt_split.json")
116 | parser.add_argument("--begin", type=int)
117 | parser.add_argument("--end", type=int)
118 | parser.add_argument("--model-name-or-path", type=str, required=True)
119 | parser.add_argument("--max-length", type=int, default=2048)
120 | args = parser.parse_args()
121 | main(args)
122 |
--------------------------------------------------------------------------------
/fastchat/data/trans_dummy.py:
--------------------------------------------------------------------------------
1 | import json
2 | from transformers import pipeline
3 | import torch
4 |
5 | translator_en2zh = None
6 |
7 |
8 | def load_model():
9 | global translator_en2zh
10 | if translator_en2zh is None:
11 | translator_en2zh = pipeline(
12 | "translation", model="Helsinki-NLP/opus-mt-en-zh", device=0)
13 | print("load model ok")
14 |
15 |
16 | def trans_en2zh(input):
17 | global translator_en2zh
18 | return translator_en2zh(input)
19 |
20 |
21 | def trans_dummy():
22 | with open('./data/dummy_en.json', 'r+', encoding='utf-8') as f:
23 | data = json.load(f)
24 | j = len(data)
25 | i = 0
26 | for items in data:
27 | i = i + 1
28 | print(str(i) + ' of ' + str(j))
29 | items['id'] = 'cn_' + items['id']
30 | for item in items["conversations"] :
31 | item["value"] = trans_en2zh(item['value'])[0]['translation_text']
32 | with open('./data/dummy_cn.json', 'w', encoding='utf-8') as f:
33 | json.dump(data, f,ensure_ascii=False,indent=2)
34 |
35 |
36 | if __name__ == "__main__":
37 | print("torch.cuda.is_available:", torch.cuda.is_available())
38 | load_model()
39 | trans_dummy()
40 |
--------------------------------------------------------------------------------
/fastchat/train/train_lora.py:
--------------------------------------------------------------------------------
1 | # Usage: deepspeed train_lora.py --deepspeed <$PATH_TO_DEEPSPEED_CONFIG>
2 |
3 | # Adopted from tatsu-lab@stanford_alpaca. Below is the original copyright:
4 | # Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
5 | #
6 | # Licensed under the Apache License, Version 2.0 (the "License");
7 | # you may not use this file except in compliance with the License.
8 | # You may obtain a copy of the License at
9 | #
10 | # http://www.apache.org/licenses/LICENSE-2.0
11 | #
12 | # Unless required by applicable law or agreed to in writing, software
13 | # distributed under the License is distributed on an "AS IS" BASIS,
14 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15 | # See the License for the specific language governing permissions and
16 | # limitations under the License.
17 | import os
18 | from dataclasses import dataclass, field
19 | import logging
20 | import pathlib
21 | import typing
22 | import torch
23 |
24 | from deepspeed import zero
25 | from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
26 | from peft import (
27 | LoraConfig,
28 | get_peft_model,
29 | get_peft_model_state_dict,
30 | prepare_model_for_kbit_training,
31 | set_peft_model_state_dict,
32 | )
33 | from transformers import LlamaForCausalLM, LlamaTokenizer
34 | import transformers
35 | from transformers import Trainer
36 |
37 | from fastchat.train.train import (
38 | DataArguments,
39 | ModelArguments,
40 | TrainingArguments,
41 | make_supervised_data_module,
42 | )
43 |
44 |
45 | @dataclass
46 | class LoraArguments:
47 | lora_r: int = 8
48 | lora_alpha: int = 16
49 | lora_dropout: float = 0.05
50 | lora_target_modules: typing.List[str] = field(
51 | default_factory=lambda: ["q_proj", "v_proj"]
52 | )
53 | lora_weight_path: str = ""
54 | bias: str = "none"
55 |
56 |
57 | def maybe_zero_3(param):
58 | if hasattr(param, "ds_id"):
59 | assert param.ds_status == ZeroParamStatus.NOT_AVAILABLE
60 | with zero.GatheredParameters([param]):
61 | param = param.data.cpu().clone().detach()
62 | return param
63 |
64 |
65 | # Borrowed from peft.utils.get_peft_model_state_dict
66 | def get_peft_state_maybe_zero_3(state_dict, bias):
67 | if bias == "none":
68 | to_return = {
69 | k: state_dict[k].cpu().clone().detach() for k in state_dict if "lora_" in k
70 | }
71 | elif bias == "all":
72 | to_return = {
73 | k: state_dict[k] for k in state_dict if "lora_" in k or "bias" in k
74 | }
75 | elif bias == "lora_only":
76 | to_return = {}
77 | for k in state_dict:
78 | if "lora_" in k:
79 | to_return[k] = state_dict[k]
80 | bias_name = k.split("lora_")[0] + "bias"
81 | if bias_name in state_dict:
82 | to_return[bias_name] = state_dict[bias_name]
83 | else:
84 | raise NotImplementedError
85 | to_return = {k: maybe_zero_3(v) for k, v in to_return.items()}
86 | return to_return
87 |
88 |
89 | def train():
90 | parser = transformers.HfArgumentParser(
91 | (ModelArguments, DataArguments, TrainingArguments, LoraArguments)
92 | )
93 | (
94 | model_args,
95 | data_args,
96 | training_args,
97 | lora_args,
98 | ) = parser.parse_args_into_dataclasses()
99 | print("开始装载模型")
100 | device_map = "auto"
101 | world_size = int(os.environ.get("WORLD_SIZE", 1))
102 | ddp = world_size != 1
103 | if ddp:
104 | device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)}
105 | model = LlamaForCausalLM.from_pretrained(
106 | model_args.model_name_or_path,
107 | cache_dir=training_args.cache_dir,
108 | load_in_8bit=True,
109 | torch_dtype=torch.float16,
110 | device_map=device_map
111 | )
112 | print("装载模型完成")
113 | model = prepare_model_for_kbit_training(model)
114 | print("模型处理为int8")
115 | lora_config = LoraConfig(
116 | r=lora_args.lora_r,
117 | lora_alpha=lora_args.lora_alpha,
118 | target_modules=lora_args.lora_target_modules,
119 | lora_dropout=lora_args.lora_dropout,
120 | bias=lora_args.bias,
121 | task_type="CAUSAL_LM",
122 | )
123 | model = get_peft_model(model, lora_config)
124 | print("模型用peft处理")
125 | if training_args.deepspeed is not None and training_args.local_rank == 0:
126 | model.print_trainable_parameters()
127 | tokenizer = transformers.AutoTokenizer.from_pretrained(
128 | model_args.model_name_or_path,
129 | cache_dir=training_args.cache_dir,
130 | model_max_length=training_args.model_max_length,
131 | padding_side="right",
132 | use_fast=False,
133 | )
134 | tokenizer.pad_token = tokenizer.unk_token
135 | print("装载tokenizer")
136 | data_module = make_supervised_data_module(
137 | tokenizer=tokenizer, data_args=data_args)
138 | if torch.cuda.device_count() > 1:
139 | model.is_parallelizable = True
140 | model.model_parallel = True
141 | model.config.use_cache = False
142 | print("装载训练数据")
143 | trainer = Trainer(
144 | model=model, tokenizer=tokenizer, args=training_args, **data_module
145 | )
146 | print("准备训练参数")
147 | # model.config.use_cache = False
148 | print("开始训练模型")
149 | if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
150 | trainer.train(resume_from_checkpoint=True)
151 | else:
152 | trainer.train()
153 | print("准备训练状态")
154 | trainer.save_state()
155 | print("保存训练模型")
156 | # Save states. Weights might be a placeholder in zero3 and need a gather
157 | state_dict = get_peft_state_maybe_zero_3(
158 | model.state_dict(), lora_args.bias)
159 | if training_args.local_rank == 0:
160 | model.save_pretrained(training_args.output_dir, state_dict=state_dict)
161 |
162 |
163 | if __name__ == "__main__":
164 | with torch.autocast("cuda"):
165 | train()
166 |
--------------------------------------------------------------------------------
/generate.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 |
4 | import argparse
5 | import torch
6 | import transformers
7 | from peft import PeftModel
8 | from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer
9 |
10 | from utils.prompter import Prompter
11 |
12 | if torch.cuda.is_available():
13 | device = "cuda"
14 | else:
15 | device = "cpu"
16 |
17 | try:
18 | if torch.backends.mps.is_available():
19 | device = "mps"
20 | except: # noqa: E722
21 | pass
22 |
23 |
24 | def main(
25 | load_8bit: bool = False,
26 | base_model: str = "",
27 | lora_weights: str = "tloen/alpaca-lora-7b",
28 | prompt_template: str = ""
29 | ):
30 | base_model = base_model or os.environ.get("BASE_MODEL", "")
31 | assert (
32 | base_model
33 | ), "Please specify a --base_model, e.g. --base_model='huggyllama/llama-7b'"
34 |
35 | prompter = Prompter(prompt_template)
36 | tokenizer = LlamaTokenizer.from_pretrained(base_model)
37 | if device == "cuda":
38 | model = LlamaForCausalLM.from_pretrained(
39 | base_model,
40 | load_in_8bit=load_8bit,
41 | torch_dtype=torch.float16,
42 | device_map="auto",
43 | )
44 | model = PeftModel.from_pretrained(
45 | model,
46 | lora_weights,
47 | torch_dtype=torch.float16,
48 | )
49 | elif device == "mps":
50 | model = LlamaForCausalLM.from_pretrained(
51 | base_model,
52 | device_map={"": device},
53 | torch_dtype=torch.float16,
54 | )
55 | model = PeftModel.from_pretrained(
56 | model,
57 | lora_weights,
58 | device_map={"": device},
59 | torch_dtype=torch.float16,
60 | )
61 | else:
62 | model = LlamaForCausalLM.from_pretrained(
63 | base_model, device_map={"": device}, low_cpu_mem_usage=True
64 | )
65 | model = PeftModel.from_pretrained(
66 | model,
67 | lora_weights,
68 | device_map={"": device},
69 | )
70 |
71 | # unwind broken decapoda-research config
72 | model.config.pad_token_id = tokenizer.pad_token_id = 0 # unk
73 | model.config.bos_token_id = 1
74 | model.config.eos_token_id = 2
75 |
76 | if not load_8bit:
77 | model.half() # seems to fix bugs for some users.
78 |
79 | model.eval()
80 | if torch.__version__ >= "2" and sys.platform != "win32":
81 | model = torch.compile(model)
82 |
83 | def evaluate(
84 | instruction,
85 | input=None,
86 | temperature=0.1,
87 | top_p=0.75,
88 | top_k=40,
89 | num_beams=4,
90 | max_new_tokens=128,
91 | stream_output=False,
92 | **kwargs,
93 | ):
94 | prompt = prompter.generate_prompt(instruction, input)
95 | inputs = tokenizer(prompt, return_tensors="pt")
96 | input_ids = inputs["input_ids"].to(device)
97 | generation_config = GenerationConfig(
98 | temperature=temperature,
99 | top_p=top_p,
100 | top_k=top_k,
101 | num_beams=num_beams,
102 | **kwargs,
103 | )
104 |
105 | generate_params = {
106 | "input_ids": input_ids,
107 | "generation_config": generation_config,
108 | "return_dict_in_generate": True,
109 | "output_scores": True,
110 | "max_new_tokens": max_new_tokens,
111 | }
112 |
113 | with torch.no_grad():
114 | generation_output = model.generate(
115 | input_ids=input_ids,
116 | generation_config=generation_config,
117 | return_dict_in_generate=True,
118 | output_scores=True,
119 | max_new_tokens=max_new_tokens,
120 | )
121 | s = generation_output.sequences[0]
122 | output = tokenizer.decode(s)
123 | return prompter.get_response(output)
124 |
125 | while True:
126 | instruction = input("Input:")
127 | if len(instruction.strip()) == 0:
128 | break
129 | print("Response:", evaluate(instruction))
130 |
131 |
132 | if __name__ == "__main__":
133 | # parse args
134 | parser = argparse.ArgumentParser()
135 | parser.add_argument('--base_model', default=None, type=str, required=True)
136 | parser.add_argument('--lora_weights', default=None, type=str,
137 | help="If None, perform inference on the base model")
138 | parser.add_argument('--load_8bit', action='store_true',
139 | help='only use CPU for inference')
140 | args = parser.parse_args()
141 | main(args.load_8bit, args.base_model, args.lora_weights, "")
142 |
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | transformers
2 | torch
3 | bs4
4 | markdownify
5 | polyglot
6 | pyicu
7 | pycld2
8 | tqdm
9 | accelerate
10 | einops
11 | flash_attn==1.0.5
12 | peft
13 | deepspeed
14 | bitsandbytes==0.37.2
15 | SentencePiece
16 | fschat==0.2.10
17 | opencc
18 |
--------------------------------------------------------------------------------
/templates/alpaca.json:
--------------------------------------------------------------------------------
1 | {
2 | "description": "Template used by Alpaca-LoRA.",
3 | "prompt_input": "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n",
4 | "prompt_no_input": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Response:\n",
5 | "response_split": "### Response:"
6 | }
7 |
--------------------------------------------------------------------------------
/utils/prompter.py:
--------------------------------------------------------------------------------
1 | """
2 | A dedicated helper to manage templates and prompt building.
3 | """
4 |
5 | import json
6 | import os.path as osp
7 | from typing import Union
8 |
9 |
10 | class Prompter(object):
11 | __slots__ = ("template", "_verbose")
12 |
13 | def __init__(self, template_name: str = "", verbose: bool = False):
14 | self._verbose = verbose
15 | if not template_name:
16 | # Enforce the default here, so the constructor can be called with '' and will not break.
17 | template_name = "alpaca"
18 | file_name = osp.join("templates", f"{template_name}.json")
19 | if not osp.exists(file_name):
20 | raise ValueError(f"Can't read {file_name}")
21 | with open(file_name) as fp:
22 | self.template = json.load(fp)
23 | if self._verbose:
24 | print(
25 | f"Using prompt template {template_name}: {self.template['description']}"
26 | )
27 |
28 | def generate_prompt(
29 | self,
30 | instruction: str,
31 | input: Union[None, str] = None,
32 | label: Union[None, str] = None,
33 | ) -> str:
34 | # returns the full prompt from instruction and optional input
35 | # if a label (=response, =output) is provided, it's also appended.
36 | if input:
37 | res = self.template["prompt_input"].format(
38 | instruction=instruction, input=input
39 | )
40 | else:
41 | res = self.template["prompt_no_input"].format(
42 | instruction=instruction
43 | )
44 | if label:
45 | res = f"{res}{label}"
46 | if self._verbose:
47 | print(res)
48 | return res
49 |
50 | def get_response(self, output: str) -> str:
51 | return output.split(self.template["response_split"])[1].strip()
52 |
--------------------------------------------------------------------------------