wehub/flagopen--flagembedding
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:39:21 +08:00
commit bc34f6df14
1149 changed files with 328099 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/README.md
+104
View File
@@ -0,0 +1,104 @@
# 1. Introduction
In this example, we show how to use scripts to make your fine-tuning process more convenient
# 2. Installation
```shell
git clone https://github.com/FlagOpen/FlagEmbedding.git
cd FlagEmbedding/scripts
```
# 3. Usage
### Hard Negatives
Hard negatives is a widely used method to improve the quality of sentence embedding. You can mine hard negatives following this command:
```shell
python hn_mine.py \
--input_file toy_finetune_data.jsonl \
--output_file toy_finetune_data_minedHN.jsonl \
--range_for_sampling 2-200 \
--negative_number 15 \
--use_gpu_for_searching \
--embedder_name_or_path BAAI/bge-base-en-v1.5
```
- **`input_file`**: json data for finetuning. This script will retrieve top-k documents for each query, and random sample negatives from the top-k documents (not including the positive documents).
- **`output_file`**: path to save JSON data with mined hard negatives for finetuning
- **`negative_number`**: the number of sampled negatives
- **`range_for_sampling`**: where to sample negative. For example, `2-100` means sampling `negative_number` negatives from top2-top200 documents. **You can set larger value to reduce the difficulty of negatives (e.g., set it `60-300` to sample negatives from top60-300 passages)**
- **`candidate_pool`**: The pool to retrieval. The default value is None, and this script will retrieve from the combination of all `neg` in `input_file`. If provided, it should be a jsonl file, each line is a dict with a key `text`. If input a candidate_pool, this script will retrieve negatives from this file.
- **`use_gpu_for_searching`**: whether to use faiss-gpu to retrieve negatives.
- **`search_batch_size`**: batch size for searching. Default is 64.
- **`embedder_name_or_path`**: The name or path to the embedder.
- **`embedder_model_class`**: Class of the model used for embedding (current options include 'encoder-only-base', 'encoder-only-m3', 'decoder-only-base', 'decoder-only-icl'.). Default is None. For the custom model, you should set this argument.
- **`normalize_embeddings`**: Set to `True` to normalize embeddings.
- **`pooling_method`**: The pooling method for the embedder.
- **`use_fp16`**: Use FP16 precision for inference.
- **`devices`**: List of devices used for inference.
- **`query_instruction_for_retrieval`**, **`query_instruction_format_for_retrieval`**: Instructions and format for query during retrieval.
- **`examples_for_task`**, **`examples_instruction_format`**: Example tasks and their instructions format. This is only used when `embedder_model_class` is set to `decoder-only-icl`.
- **`trust_remote_code`**: Set to `True` to trust remote code execution.
- **`cache_dir`**: Cache directory for models.
- **`embedder_batch_size`**: Batch sizes for embedding and reranking.
- **`embedder_query_max_length`**, **`embedder_passage_max_length`**: Maximum length for embedding queries and passages.
### Teacher Scores
Teacher scores can be used for model distillation. You can obtain the scores using the following command:
```shell
python add_reranker_score.py \
--input_file toy_finetune_data_minedHN.jsonl \
--output_file toy_finetune_data_score.jsonl \
--reranker_name_or_path BAAI/bge-reranker-v2-m3
```
- **`input_file`**: path to save JSON data with mined hard negatives for finetuning
- **`output_file`**: path to save JSON data with scores for finetuning
- **`use_fp16`**: Whether to use fp16 for inference. Default: True
- **`devices`**: Devices to use for inference. Default: None, multiple values allowed
- **`trust_remote_code`**: Trust remote code. Default: False
- **`reranker_name_or_path`**: The reranker name or path. Default: None
- **`reranker_model_class`**: The reranker model class. Available classes: ['auto', 'encoder-only-base', 'decoder-only-base', 'decoder-only-layerwise', 'decoder-only-lightweight']. Default: auto
- **`reranker_peft_path`**: The reranker peft path. Default: None
- **`use_bf16`**: Whether to use bf16 for inference. Default: False
- **`query_instruction_for_rerank`**: Instruction for query. Default: None
- **`query_instruction_format_for_rerank`**: Format for query instruction. Default: {{}{}}
- **`passage_instruction_for_rerank`**: Instruction for passage. Default: None
- **`passage_instruction_format_for_rerank`**: Format for passage instruction. Default: {{}{}}
- **`cache_dir`**: Cache directory for models. Default: None
- **`reranker_batch_size`**: Batch size for inference. Default: 3000
- **`reranker_query_max_length`**: Max length for reranking queries. Default: None
- **`reranker_max_length`**: Max length for reranking. Default: 512
- **`normalize`**: Whether to normalize the reranking scores. Default: False
- **`prompt`**: The prompt for the reranker. Default: None
- **`cutoff_layers`**: The output layers of layerwise/lightweight reranker. Default: None
- **`compress_ratio`**: The compress ratio of lightweight reranker. Default: 1
- **`compress_layers`**: The compress layers of lightweight reranker. Default: None, multiple values allowed
### Split Data by Length
You can split the data using the following command:
```shell
python split_data_by_length.py \
--input_path train_data \
--output_dir train_data_split \
--cache_dir .cache \
--log_name .split_log \
--length_list 0 500 1000 2000 3000 4000 5000 6000 7000 \
--model_name_or_path BAAI/bge-m3 \
--num_proc 16
```
- **`input_path`**: The path of input data. It can be a file or a directory containing multiple files.
- **`output_dir`**: The directory of output data. The split data files will be saved to this directory.
- **`cache_dir`**: The cache directory. Default: None
- **`log_name`**: The name of the log file. Default: `.split_log`, which will be saved to `output_dir`
- **`length_list`**: The length list to split. Default: [0, 500, 1000, 2000, 3000, 4000, 5000, 6000, 7000]
- **`model_name_or_path`**: The model name or path of the tokenizer. Default: `BAAI/bge-m3`
- **`num_proc`**: The number of processes. Default: 16
- **`overwrite`**: Whether to overwrite the output file. Default: False
scripts/add_reranker_score.py
+144
View File
@@ -0,0 +1,144 @@
import json
from typing import Optional, List
from dataclasses import dataclass, field
from transformers import HfArgumentParser
from FlagEmbedding import FlagAutoReranker
@dataclass
class ScoreArgs:
input_file: str = field(
default=None, metadata={"help": "The input jsonl file, each line includes query, pos and neg."}
)
output_file: str = field(
default=None, metadata={"help": "The output jsonl file, it includes query, pos, neg, pos_scores and neg_scores."}
)
@dataclass
class ModelArgs:
use_fp16: bool = field(
default=True, metadata={"help": "whether to use fp16 for inference"}
)
devices: Optional[str] = field(
default=None, metadata={"help": "Devices to use for inference.", "nargs": "+"}
)
trust_remote_code: bool = field(
default=False, metadata={"help": "Trust remote code"}
)
reranker_name_or_path: Optional[str] = field(
default=None, metadata={"help": "The reranker name or path."}
)
reranker_model_class: Optional[str] = field(
default=None, metadata={"help": "The reranker model class. Available classes: ['encoder-only-base', 'decoder-only-base', 'decoder-only-layerwise', 'decoder-only-lightweight']. Default: None. For the custom model, you need to specify the model class.", "choices": ["encoder-only-base", "decoder-only-base", "decoder-only-layerwise", "decoder-only-lightweight"]}
)
reranker_peft_path: Optional[str] = field(
default=None, metadata={"help": "The reranker peft path."}
)
use_bf16: bool = field(
default=False, metadata={"help": "whether to use bf16 for inference"}
)
query_instruction_for_rerank: Optional[str] = field(
default=None, metadata={"help": "Instruction for query"}
)
query_instruction_format_for_rerank: str = field(
default="{}{}", metadata={"help": "Format for query instruction"}
)
passage_instruction_for_rerank: Optional[str] = field(
default=None, metadata={"help": "Instruction for passage"}
)
passage_instruction_format_for_rerank: str = field(
default="{}{}", metadata={"help": "Format for passage instruction"}
)
cache_dir: str = field(
default=None, metadata={"help": "Cache directory for models."}
)
# ================ for inference ===============
reranker_batch_size: int = field(
default=3000, metadata={"help": "Batch size for inference."}
)
reranker_query_max_length: Optional[int] = field(
default=None, metadata={"help": "Max length for reranking."}
)
reranker_max_length: int = field(
default=512, metadata={"help": "Max length for reranking."}
)
normalize: bool = field(
default=False, metadata={"help": "whether to normalize the reranking scores"}
)
prompt: Optional[str] = field(
default=None, metadata={"help": "The prompt for the reranker."}
)
cutoff_layers: List[int] = field(
default=None, metadata={"help": "The output layers of layerwise/lightweight reranker."}
)
compress_ratio: int = field(
default=1, metadata={"help": "The compress ratio of lightweight reranker."}
)
compress_layers: Optional[int] = field(
default=None, metadata={"help": "The compress layers of lightweight reranker.", "nargs": "+"}
)
def main(score_args: ScoreArgs, model_args: ModelArgs):
reranker = FlagAutoReranker.from_finetuned(
model_name_or_path=model_args.reranker_name_or_path,
model_class=model_args.reranker_model_class,
peft_path=model_args.reranker_peft_path,
use_fp16=model_args.use_fp16,
use_bf16=model_args.use_bf16,
query_instruction_for_rerank=model_args.query_instruction_for_rerank,
query_instruction_format=model_args.query_instruction_format_for_rerank,
passage_instruction_for_rerank=model_args.passage_instruction_for_rerank,
passage_instruction_format=model_args.passage_instruction_format_for_rerank,
cache_dir=model_args.cache_dir,
trust_remote_code=model_args.trust_remote_code,
devices=model_args.devices,
normalize=model_args.normalize,
prompt=model_args.prompt,
cutoff_layers=model_args.cutoff_layers,
compress_layers=model_args.compress_layers,
compress_ratio=model_args.compress_ratio,
batch_size=model_args.reranker_batch_size,
query_max_length=model_args.reranker_query_max_length,
max_length=model_args.reranker_max_length,
)
pairs = []
data = []
with open(score_args.input_file) as f:
for line in f:
data.append(json.loads(line))
for p in data[-1]['pos']:
pairs.append((data[-1]['query'], p))
for p in data[-1]['neg']:
pairs.append((data[-1]['query'], p))
scores = reranker.compute_score(pairs)
score_idx = 0
for i in range(len(data)):
data[i]['pos_scores'] = []
data[i]['neg_scores'] = []
for _ in range(len(data[i]['pos'])):
data[i]['pos_scores'].append(float(scores[score_idx]))
score_idx += 1
for _ in range(len(data[i]['neg'])):
data[i]['neg_scores'].append(float(scores[score_idx]))
score_idx += 1
with open(score_args.output_file, 'w') as f:
for d in data:
f.write(json.dumps(d) + '\n')
if __name__ == "__main__":
parser = HfArgumentParser((
ScoreArgs,
ModelArgs
))
score_args, model_args = parser.parse_args_into_dataclasses()
score_args: ScoreArgs
model_args: ModelArgs
main(score_args, model_args)
scripts/hn_mine.py
+246
View File
@@ -0,0 +1,246 @@
import json
import random
import numpy as np
from tqdm import tqdm
from typing import Optional
from dataclasses import dataclass, field
import faiss
from transformers import HfArgumentParser
from FlagEmbedding import FlagAutoModel
from FlagEmbedding.abc.inference import AbsEmbedder
@dataclass
class DataArgs:
"""
Data arguments for hard negative mining.
"""
input_file: str = field(
metadata={"help": "The input file for hard negative mining."}
)
output_file: str = field(
metadata={"help": "The output file for hard negative mining."}
)
candidate_pool: Optional[str] = field(
default=None, metadata={"help": "The candidate pool for hard negative mining. If provided, it should be a jsonl file, each line is a dict with a key 'text'."}
)
range_for_sampling: str = field(
default="10-210", metadata={"help": "The range to sample negatives."}
)
negative_number: int = field(
default=15, metadata={"help": "The number of negatives."}
)
use_gpu_for_searching: bool = field(
default=False, metadata={"help": "Whether to use faiss-gpu for searching."}
)
search_batch_size: int = field(
default=64, metadata={"help": "The batch size for searching."}
)
@dataclass
class ModelArgs:
"""
Model arguments for embedder.
"""
embedder_name_or_path: str = field(
metadata={"help": "The embedder name or path.", "required": True}
)
embedder_model_class: Optional[str] = field(
default=None, metadata={"help": "The embedder model class. Available classes: ['encoder-only-base', 'encoder-only-m3', 'decoder-only-base', 'decoder-only-icl']. Default: None. For the custom model, you need to specifiy the model class.", "choices": ["encoder-only-base", "encoder-only-m3", "decoder-only-base", "decoder-only-icl"]}
)
normalize_embeddings: bool = field(
default=True, metadata={"help": "whether to normalize the embeddings"}
)
pooling_method: str = field(
default="cls", metadata={"help": "The pooling method fot the embedder."}
)
use_fp16: bool = field(
default=True, metadata={"help": "whether to use fp16 for inference"}
)
devices: Optional[str] = field(
default=None, metadata={"help": "Devices to use for inference.", "nargs": "+"}
)
query_instruction_for_retrieval: Optional[str] = field(
default=None, metadata={"help": "Instruction for query"}
)
query_instruction_format_for_retrieval: str = field(
default="{}{}", metadata={"help": "Format for query instruction"}
)
examples_for_task: Optional[str] = field(
default=None, metadata={"help": "Examples for task"}
)
examples_instruction_format: str = field(
default="{}{}", metadata={"help": "Format for examples instruction"}
)
trust_remote_code: bool = field(
default=False, metadata={"help": "Trust remote code"}
)
cache_dir: str = field(
default=None, metadata={"help": "Cache directory for models."}
)
# ================ for inference ===============
batch_size: int = field(
default=3000, metadata={"help": "Batch size for inference."}
)
embedder_query_max_length: int = field(
default=512, metadata={"help": "Max length for query."}
)
embedder_passage_max_length: int = field(
default=512, metadata={"help": "Max length for passage."}
)
def __post_init__(self):
# replace "\\n" with "\n"
if "\\n" in self.query_instruction_format_for_retrieval:
self.query_instruction_format_for_retrieval = self.query_instruction_format_for_retrieval.replace("\\n", "\n")
if "\\n" in self.examples_instruction_format:
self.examples_instruction_format = self.examples_instruction_format.replace("\\n", "\n")
def create_index(embeddings: np.ndarray, use_gpu: bool = False):
index = faiss.IndexFlatIP(len(embeddings[0]))
embeddings = np.asarray(embeddings, dtype=np.float32)
if use_gpu:
co = faiss.GpuMultipleClonerOptions()
co.shard = True
co.useFloat16 = True
index = faiss.index_cpu_to_all_gpus(index, co=co)
index.add(embeddings)
return index
def batch_search(
index: faiss.Index,
query: np.ndarray,
topk: int = 200,
batch_size: int = 64
):
all_scores, all_inxs = [], []
for start_index in tqdm(range(0, len(query), batch_size), desc="Batches", disable=len(query) < 256):
batch_query = query[start_index:start_index + batch_size]
batch_scores, batch_inxs = index.search(np.asarray(batch_query, dtype=np.float32), k=topk)
all_scores.extend(batch_scores.tolist())
all_inxs.extend(batch_inxs.tolist())
return all_scores, all_inxs
def get_corpus(candidate_pool: str):
corpus = []
with open(candidate_pool, "r", encoding="utf-8") as f:
for line in f.readlines():
line = json.loads(line.strip())
corpus.append(line['text'])
return corpus
def find_knn_neg(
model: AbsEmbedder,
input_file: str,
output_file: str,
candidate_pool: Optional[str] = None,
sample_range: str = "10-210",
negative_number: int = 15,
use_gpu: bool = False
):
corpus = []
queries = []
train_data = []
for line in open(input_file):
line = json.loads(line.strip())
train_data.append(line)
corpus.extend(line['pos'])
if 'neg' in line:
corpus.extend(line['neg'])
queries.append(line['query'])
if candidate_pool is not None:
if not isinstance(candidate_pool, list):
candidate_pool = get_corpus(candidate_pool)
corpus = list(set(candidate_pool))
else:
corpus = list(set(corpus))
print(f'inferencing embedding for corpus (number={len(corpus)})--------------')
p_vecs = model.encode(corpus)
print(f'inferencing embedding for queries (number={len(queries)})--------------')
q_vecs = model.encode_queries(queries)
# check if the embeddings are in dictionary format: M3Embedder
if isinstance(p_vecs, dict):
p_vecs = p_vecs["dense_vecs"]
if isinstance(q_vecs, dict):
q_vecs = q_vecs["dense_vecs"]
print('create index and search------------------')
index = create_index(p_vecs, use_gpu=use_gpu)
_, all_inxs = batch_search(index, q_vecs, topk=sample_range[-1])
assert len(all_inxs) == len(train_data)
for i, data in enumerate(train_data):
query = data['query']
inxs = all_inxs[i][sample_range[0]:sample_range[1]]
filtered_inx = []
for inx in inxs:
if inx == -1: break
if corpus[inx] not in data['pos'] and corpus[inx] != query:
filtered_inx.append(inx)
if len(filtered_inx) > negative_number:
filtered_inx = random.sample(filtered_inx, negative_number)
data['neg'] = [corpus[inx] for inx in filtered_inx]
with open(output_file, 'w') as f:
for data in train_data:
if len(data['neg']) < negative_number:
samples = random.sample(corpus, negative_number - len(data['neg']) + len(data['pos']))
samples = [sent for sent in samples if sent not in data['pos']]
data['neg'].extend(samples[: negative_number - len(data['neg'])])
f.write(json.dumps(data, ensure_ascii=False) + '\n')
def load_model(model_args: ModelArgs):
model = FlagAutoModel.from_finetuned(
model_name_or_path=model_args.embedder_name_or_path,
model_class=model_args.embedder_model_class,
normalize_embeddings=model_args.normalize_embeddings,
pooling_method=model_args.pooling_method,
use_fp16=model_args.use_fp16,
query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
query_instruction_format=model_args.query_instruction_format_for_retrieval,
devices=model_args.devices,
examples_for_task=model_args.examples_for_task,
examples_instruction_format=model_args.examples_instruction_format,
trust_remote_code=model_args.trust_remote_code,
cache_dir=model_args.cache_dir,
batch_size=model_args.batch_size,
query_max_length=model_args.embedder_query_max_length,
passage_max_length=model_args.embedder_passage_max_length,
)
return model
def main(data_args: DataArgs, model_args: ModelArgs):
model = load_model(model_args)
find_knn_neg(
model=model,
input_file=data_args.input_file,
output_file=data_args.output_file,
candidate_pool=data_args.candidate_pool,
sample_range=[int(x) for x in data_args.range_for_sampling.split('-')],
negative_number=data_args.negative_number,
use_gpu=data_args.use_gpu_for_searching
)
if __name__ == "__main__":
parser = HfArgumentParser((
DataArgs,
ModelArgs
))
data_args, model_args = parser.parse_args_into_dataclasses()
data_args: DataArgs
model_args: ModelArgs
main(data_args, model_args)
scripts/split_data_by_length.py
+213
View File
@@ -0,0 +1,213 @@
"""
python split_data_by_length.py \
--input_path train_data \
--output_dir train_data_split \
--cache_dir .cache \
--log_name .split_log \
--length_list 0 500 1000 2000 3000 4000 5000 6000 7000 \
--model_name_or_path BAAI/bge-m3 \
--num_proc 16 \
--overwrite False
"""
import os
import json
import math
import time
import argparse
import datasets
from tqdm import tqdm
from pprint import pprint
from transformers import AutoTokenizer
from datasets import load_dataset, Features, Value, Sequence
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument('--input_path', type=str, required=True, help='the path of input datas')
parser.add_argument('--output_dir', type=str, required=True, help='the dir of output datas')
parser.add_argument('--cache_dir', type=str, default=None, help='the cache dir')
parser.add_argument('--log_name', type=str, default='.split_log', help='the name of log file, default: `.split_log`, which will be saved to `output_dir`')
parser.add_argument('--length_list', type=int, default=[0, 500, 1000, 2000, 3000, 4000, 5000, 6000, 7000], nargs='+', help='the length list to split')
parser.add_argument('--model_name_or_path', type=str, default='BAAI/bge-m3', help='the model name or path of the tokenizer')
parser.add_argument('--num_proc', type=int, default=16, help='the number of process, default: 16')
parser.add_argument('--overwrite', action='store_true', default=False, help='whether to overwrite the output file, default: False')
args = parser.parse_args()
return args
class SplitByLengthHandler:
def __init__(self,
model_name_or_path: str,
cache_dir: str=None,
num_proc: int=16,
length_list: list=[0, 500, 1000, 2000, 3000, 4000, 5000, 6000, 7000],
overwrite: bool=False):
self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
self.cache_dir = cache_dir
self.num_proc = num_proc
self.length_ranges_list = self._get_length_ranges_list(length_list)
self.overwrite = overwrite
pprint(self.length_ranges_list)
def _map_func(examples):
results = {}
results['idx'] = []
results['max_length'] = []
for i in range(len(examples['query'])):
idx = examples['idx'][i]
query = examples['query'][i]
pos, neg = examples['pos'][i], examples['neg'][i]
all_texts = [query] + pos + neg
max_len = 0
for x in all_texts:
tokenized_x = self.tokenizer(x)['input_ids']
if len(tokenized_x) > max_len:
max_len = len(tokenized_x)
results['idx'].append(idx)
results['max_length'].append(max_len)
return results
self._map_func = _map_func
@staticmethod
def _get_length_ranges_list(length_list: list):
length_ranges_list = []
length_list = sorted(length_list)
for i in range(len(length_list)):
length_l = length_list[i]
if i == len(length_list) - 1:
length_r = math.inf
else:
length_r = length_list[i + 1]
assert 0 <= length_l < length_r
length_ranges_list.append((length_l, length_r))
return length_ranges_list
def _process_dir(self, dir_path: str, output_dir: str):
assert os.path.isdir(dir_path)
log_info_list = []
for file in tqdm(os.listdir(dir_path), desc=f'processing {dir_path}'):
file_path = os.path.join(dir_path, file)
if not file_path.endswith('.jsonl'):
print(f"skip {file_path} ...")
continue
output_path = os.path.join(output_dir, '.'.join(file.split('.')[:-1]))
log_info = self._process_file(file_path, output_path)
log_info_list.append(log_info)
return log_info_list
def _process_file(self, file_path: str, output_path: str):
assert not os.path.isdir(file_path)
start_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
features = Features({
'query': Value('string'),
'pos': Sequence(Value('string')),
'neg': Sequence(Value('string'))
})
kd_features = Features({
'query': Value('string'),
'pos': Sequence(Value('string')),
'neg': Sequence(Value('string')),
'pos_scores': Sequence(Value('float')),
'neg_scores': Sequence(Value('float'))
})
try:
dataset = load_dataset('json', data_files=file_path, cache_dir=self.cache_dir, features=features)['train']
except:
dataset = load_dataset('json', data_files=file_path, cache_dir=self.cache_dir, features=kd_features)['train']
dataset_with_idx_list = []
for i, data in enumerate(dataset):
data['idx'] = i
dataset_with_idx_list.append(data)
dataset_with_idx = datasets.Dataset.from_list(dataset_with_idx_list)
mapped_dataset = dataset_with_idx.map(self._map_func, batched=True, num_proc=self.num_proc)
split_info_dict = {}
for length_l, length_r in self.length_ranges_list:
save_path = output_path + f'_len-{length_l}-{length_r}.jsonl'
if os.path.exists(save_path) and not self.overwrite:
print(f'{save_path} exists, skip')
continue
idxs = mapped_dataset.filter(lambda x: length_l <= x['max_length'] < length_r, num_proc=self.num_proc)
split_dataset = dataset_with_idx.select(idxs['idx'])
split_dataset = split_dataset.remove_columns('idx')
split_info_dict[f'len-{length_l}-{length_r}'] = len(split_dataset)
if len(split_dataset) > 0:
split_dataset.to_json(save_path, force_ascii=False)
end_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
size = len(dataset)
avg_length = sum(mapped_dataset['max_length']) / size
log_info = {
'file_name': os.path.basename(file_path),
'size': size,
'avg_length': avg_length,
'file_path': file_path,
'start_time': start_time,
'end_time': end_time,
'split_info': split_info_dict
}
return log_info
def run(self, input_path: str, output_dir: str, log_name: str=None):
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if log_name is None:
log_path = os.path.join(output_dir, '.split_log')
else:
log_path = os.path.join(output_dir, log_name)
log_info_list = []
if os.path.isdir(input_path):
log_info_list = self._process_dir(input_path, output_dir)
else:
file_name = os.path.basename(input_path)
output_path = os.path.join(output_dir, '.'.join(file_name.split('.')[:-1]))
log_info = self._process_file(input_path, output_path)
log_info_list.append(log_info)
with open(log_path, 'a', encoding='utf-8') as f:
for log_info in log_info_list:
json.dump(log_info, f, ensure_ascii=False)
f.write('\n')
def main(args):
input_path = args.input_path
output_dir = args.output_dir
log_name = args.log_name
handler = SplitByLengthHandler(
model_name_or_path=args.model_name_or_path,
cache_dir=args.cache_dir,
num_proc=args.num_proc,
length_list=args.length_list if isinstance(args.length_list, list) else [args.length_list],
overwrite=args.overwrite
)
handler.run(
input_path=input_path,
output_dir=output_dir,
log_name=log_name
)
print('\nDONE!')
if __name__ == "__main__":
args = get_args()
main(args)