mirror of
https://github.com/deepset-ai/haystack.git
synced 2025-07-25 18:00:28 +00:00
834f8c4902
* Change return types of file converters * Change return types of preprocessor * Change return types of crawler * Adapt utils to functions to new return types * Adapt __init__.py to new method names * Prevent circular imports * Update Documentation & Code Style * Let DocStores' run method accept Documents * Adapt tests to new return types * Update Documentation & Code Style * Put "# type: ignore" to right place * Remove id_hash_keys property from Document primitive * Update Documentation & Code Style * Adapt tests to new return types and missing id_hash_keys property * Fix mypy * Fix mypy * Adapt PDFToTextOCRConverter * Remove id_hash_keys from RestAPI tests * Update Documentation & Code Style * Rename tests * Remove redundant setting of content_type="text" * Add DeprecationWarning * Add id_hash_keys to elasticsearch_index_to_document_store * Change document type from dict to Docuemnt in PreProcessor test * Fix file path in Tutorial 5 * Remove added output in Tutorial 5 * Update Documentation & Code Style * Fix file_paths in Tutorial 9 + fix gz files in fetch_archive_from_http * Adapt tutorials to new return types * Adapt tutorial 14 to new return types * Update Documentation & Code Style * Change assertions to HaystackErrors * Import HaystackError correctly Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
86 lines
2.6 KiB
Python
86 lines
2.6 KiB
Python
def tutorial9_dpr_training():
|
|
# Training Your Own "Dense Passage Retrieval" Model
|
|
|
|
# Here are some imports that we'll need
|
|
|
|
from haystack.nodes import DensePassageRetriever
|
|
from haystack.utils import fetch_archive_from_http
|
|
from haystack.document_stores import InMemoryDocumentStore
|
|
|
|
# Download original DPR data
|
|
# WARNING: the train set is 7.4GB and the dev set is 800MB
|
|
|
|
doc_dir = "data/tutorial9"
|
|
|
|
s3_url_train = "https://dl.fbaipublicfiles.com/dpr/data/retriever/biencoder-nq-train.json.gz"
|
|
s3_url_dev = "https://dl.fbaipublicfiles.com/dpr/data/retriever/biencoder-nq-dev.json.gz"
|
|
|
|
fetch_archive_from_http(s3_url_train, output_dir=doc_dir + "/train")
|
|
fetch_archive_from_http(s3_url_dev, output_dir=doc_dir + "/dev")
|
|
|
|
## Option 1: Training DPR from Scratch
|
|
|
|
# Here are the variables to specify our training data, the models that we use to initialize DPR
|
|
# and the directory where we'll be saving the model
|
|
|
|
train_filename = "train/biencoder-nq-train.json"
|
|
dev_filename = "dev/biencoder-nq-dev.json"
|
|
|
|
query_model = "bert-base-uncased"
|
|
passage_model = "bert-base-uncased"
|
|
|
|
save_dir = "../saved_models/dpr"
|
|
|
|
# ## Option 2: Finetuning DPR
|
|
#
|
|
# # Here are the variables you might want to use instead of the set above
|
|
# # in order to perform pretraining
|
|
#
|
|
# doc_dir = "PATH_TO_YOUR_DATA_DIR"
|
|
# train_filename = "TRAIN_FILENAME"
|
|
# dev_filename = "DEV_FILENAME"
|
|
#
|
|
# query_model = "facebook/dpr-question_encoder-single-nq-base"
|
|
# passage_model = "facebook/dpr-ctx_encoder-single-nq-base"
|
|
#
|
|
# save_dir = "..saved_models/dpr"
|
|
|
|
## Initialize DPR model
|
|
|
|
retriever = DensePassageRetriever(
|
|
document_store=InMemoryDocumentStore(),
|
|
query_embedding_model=query_model,
|
|
passage_embedding_model=passage_model,
|
|
max_seq_len_query=64,
|
|
max_seq_len_passage=256,
|
|
)
|
|
|
|
# Start training our model and save it when it is finished
|
|
|
|
retriever.train(
|
|
data_dir=doc_dir,
|
|
train_filename=train_filename,
|
|
dev_filename=dev_filename,
|
|
test_filename=dev_filename,
|
|
n_epochs=1,
|
|
batch_size=16,
|
|
grad_acc_steps=8,
|
|
save_dir=save_dir,
|
|
evaluate_every=3000,
|
|
embed_title=True,
|
|
num_positives=1,
|
|
num_hard_negatives=1,
|
|
)
|
|
|
|
## Loading
|
|
|
|
reloaded_retriever = DensePassageRetriever.load(load_dir=save_dir, document_store=None)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tutorial9_dpr_training()
|
|
|
|
# This Haystack script was made with love by deepset in Berlin, Germany
|
|
# Haystack: https://github.com/deepset-ai/haystack
|
|
# deepset: https://deepset.ai/
|