| |
| """ |
| Build Faiss Index for all document in ./data |
| """ |
|
|
| import argparse |
| from build_index.base import FileReader, Document |
| from build_index.process import split_documents |
| from build_index.pricing import check_price |
| from build_index.doc2vec import doc2vec |
|
|
|
|
| def process_directory(directory, folder_name, model, max_tokens): |
| reader = FileReader(directory, folder_name) |
| docs = reader.load_data() |
| print(f'Before chunking {len(docs)} docs') |
| new_docs = split_documents(documents=docs, max_tokens=max_tokens) |
| print(f'After chunking {len(new_docs)} docs') |
| docs = [Document.to_langchain_format(doc) for doc in new_docs] |
|
|
| doc2vec(docs, model, folder_name) |
| return docs |
|
|
|
|
| def process_files(files, model, max_tokens): |
| reader = FileReader(input_files=files) |
| docs = reader.load_data() |
| print(f'Before chunking {len(docs)} docs') |
| new_docs = split_documents(documents=docs, max_tokens=max_tokens) |
| print(f'After chunking {len(new_docs)} docs') |
| docs = [Document.to_langchain_format(doc) for doc in new_docs] |
|
|
| doc2vec(docs, model) |
| return docs |
|
|
|
|
| if __name__ =='__main__': |
| import argparse |
| parser = argparse.ArgumentParser() |
| parser.add_argument("directory", type=str, deafult='./data') |
| parser.add_argument("folder", type=str, default='半导体') |
| parser.add_argument("model", type=str, default='mpnet') |
| parser.add_argument('max_tokens', type=int, default=514) |
| args = parser.parse_args() |
|
|
| process_directory(args.directory, args.folder, args.model, args.max_token) |
|
|
| process_directory('./data', '半导体', 'mpnet', 514) |
| process_directory('./data', '数字经济', 'cohere', 514) |
| process_directory('./data', '数字经济', 'cohere', 514) |
| process_directory('./data', '两会', 'cohere', 514) |
|
|
| from glob import glob |
| files = glob('./data/两会/*.pdf') |
| files = files[:3] |
| process_files(files[:3], 'cohere', 514) |
