CarlosMalaga's picture
Update relik/inference/annotator.py
0f8a590 verified
raw
history blame
37 kB
import logging
import os
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Union
import hydra
import torch
from omegaconf import DictConfig, OmegaConf
from relik.inference.data.splitters.blank_sentence_splitter import BlankSentenceSplitter
from relik.common.log import get_logger
from relik.common.upload import get_logged_in_username, upload
from relik.common.utils import CONFIG_NAME, from_cache
from relik.inference.data.objects import (
AnnotationType,
RelikOutput,
Span,
TaskType,
Triples,
)
from relik.inference.data.splitters.base_sentence_splitter import BaseSentenceSplitter
from relik.inference.data.splitters.spacy_sentence_splitter import SpacySentenceSplitter
from relik.inference.data.splitters.window_based_splitter import WindowSentenceSplitter
from relik.inference.data.tokenizers.spacy_tokenizer import SpacyTokenizer
from relik.inference.data.window.manager import WindowManager
from relik.reader.data.relik_reader_sample import RelikReaderSample
from relik.reader.pytorch_modules.base import RelikReaderBase
from relik.reader.pytorch_modules.span import RelikReaderForSpanExtraction
from relik.reader.pytorch_modules.triplet import RelikReaderForTripletExtraction
from relik.retriever.indexers.base import BaseDocumentIndex
from relik.retriever.indexers.document import Document
from relik.retriever.pytorch_modules import PRECISION_MAP
from relik.retriever.pytorch_modules.model import GoldenRetriever
# set tokenizers parallelism to False
os.environ["TOKENIZERS_PARALLELISM"] = os.getenv("TOKENIZERS_PARALLELISM", "false")
LOG_QUERY = os.getenv("RELIK_LOG_QUERY_ON_FILE", "false").lower() == "true"
logger = get_logger(__name__, level=logging.INFO)
file_logger = None
if LOG_QUERY:
RELIK_LOG_PATH = Path(__file__).parent.parent.parent / "relik.log"
# create file handler which logs even debug messages
fh = logging.FileHandler(RELIK_LOG_PATH)
fh.setLevel(logging.INFO)
file_logger = get_logger("relik", level=logging.INFO)
file_logger.addHandler(fh)
class Relik:
"""
Relik main class. It is a wrapper around a retriever and a reader.
Args:
retriever (:obj:`GoldenRetriever`):
The retriever to use.
reader (:obj:`RelikReaderBase`):
The reader to use.
document_index (:obj:`BaseDocumentIndex`, `optional`):
The document index to use. If `None`, the retriever's document index will be used.
device (`str`, `optional`, defaults to `cpu`):
The device to use for both the retriever and the reader.
retriever_device (`str`, `optional`, defaults to `None`):
The device to use for the retriever. If `None`, the `device` argument will be used.
document_index_device (`str`, `optional`, defaults to `None`):
The device to use for the document index. If `None`, the `device` argument will be used.
reader_device (`str`, `optional`, defaults to `None`):
The device to use for the reader. If `None`, the `device` argument will be used.
precision (`int`, `str` or `torch.dtype`, `optional`, defaults to `32`):
The precision to use for both the retriever and the reader.
retriever_precision (`int`, `str` or `torch.dtype`, `optional`, defaults to `None`):
The precision to use for the retriever. If `None`, the `precision` argument will be used.
document_index_precision (`int`, `str` or `torch.dtype`, `optional`, defaults to `None`):
The precision to use for the document index. If `None`, the `precision` argument will be used.
reader_precision (`int`, `str` or `torch.dtype`, `optional`, defaults to `None`):
The precision to use for the reader. If `None`, the `precision` argument will be used.
metadata_fields (`list[str]`, `optional`, defaults to `None`):
The fields to add to the candidates for the reader.
top_k (`int`, `optional`, defaults to `None`):
The number of candidates to retrieve for each window.
window_size (`int`, `optional`, defaults to `None`):
The size of the window. If `None`, the whole text will be annotated.
window_stride (`int`, `optional`, defaults to `None`):
The stride of the window. If `None`, there will be no overlap between windows.
**kwargs:
Additional keyword arguments to pass to the retriever and the reader.
"""
def __init__(
self,
retriever: GoldenRetriever | DictConfig | Dict | None = None,
reader: RelikReaderBase | DictConfig | None = None,
device: str | None = None,
retriever_device: str | None = None,
document_index_device: str | None = None,
reader_device: str | None = None,
precision: int | str | torch.dtype | None = None,
retriever_precision: int | str | torch.dtype | None = None,
document_index_precision: int | str | torch.dtype | None = None,
reader_precision: int | str | torch.dtype | None = None,
task: TaskType | str = TaskType.SPAN,
metadata_fields: list[str] | None = None,
top_k: int | None = None,
window_size: int | str | None = None,
window_stride: int | None = None,
retriever_kwargs: Dict[str, Any] | None = None,
reader_kwargs: Dict[str, Any] | None = None,
**kwargs,
) -> None:
# parse task into a TaskType
if isinstance(task, str):
try:
task = TaskType(task.lower())
except ValueError:
raise ValueError(
f"Task `{task}` not recognized. "
f"Please choose one of {list(TaskType)}."
)
self.task = task
# organize devices
if device is not None:
if retriever_device is None:
retriever_device = device
if document_index_device is None:
document_index_device = device
if reader_device is None:
reader_device = device
# organize precision
if precision is not None:
if retriever_precision is None:
retriever_precision = precision
if document_index_precision is None:
document_index_precision = precision
if reader_precision is None:
reader_precision = precision
# retriever
self.retriever: Dict[TaskType, GoldenRetriever] = {
TaskType.SPAN: None,
TaskType.TRIPLET: None,
}
if retriever:
# check retriever type, it can be a GoldenRetriever, a DictConfig or a Dict
if not isinstance(retriever, (GoldenRetriever, DictConfig, Dict)):
raise ValueError(
f"`retriever` must be a `GoldenRetriever`, a `DictConfig` or "
f"a `Dict`, got `{type(retriever)}`."
)
# we need to check weather the DictConfig is a DictConfig for an instance of GoldenRetriever
# or a primitive Dict
if isinstance(retriever, DictConfig):
# then it is probably a primitive Dict
if "_target_" not in retriever:
retriever = OmegaConf.to_container(retriever, resolve=True)
# convert the key to TaskType
try:
retriever = {
TaskType(k.lower()): v for k, v in retriever.items()
}
except ValueError as e:
raise ValueError(
f"Please choose a valid task type (one of {list(TaskType)}) for each retriever."
) from e
if isinstance(retriever, Dict):
# convert the key to TaskType
retriever = {TaskType(k): v for k, v in retriever.items()}
else:
retriever = {task: retriever}
# instantiate each retriever
if self.task in [TaskType.SPAN, TaskType.BOTH]:
self.retriever[TaskType.SPAN] = self._instantiate_retriever(
retriever[TaskType.SPAN],
retriever_device,
retriever_precision,
None,
document_index_device,
document_index_precision,
)
if self.task in [TaskType.TRIPLET, TaskType.BOTH]:
self.retriever[TaskType.TRIPLET] = self._instantiate_retriever(
retriever[TaskType.TRIPLET],
retriever_device,
retriever_precision,
None,
document_index_device,
document_index_precision,
)
# clean up None retrievers from the dictionary
self.retriever = {
task_type: r for task_type, r in self.retriever.items() if r is not None
}
# torch compile
# self.retriever = {task_type: torch.compile(r, backend="onnxrt") for task_type, r in self.retriever.items()}
# reader
self.reader: RelikReaderBase | None = None
if reader:
reader = (
hydra.utils.instantiate(
reader,
device=reader_device,
precision=reader_precision,
)
if isinstance(reader, DictConfig)
else reader
)
reader.training = False
reader.eval()
if reader_device is not None:
logger.info(f"Moving reader to `{reader_device}`.")
reader.to(reader_device)
if reader_precision is not None and reader.precision != PRECISION_MAP[reader_precision]:
logger.info(
f"Setting precision of reader to `{PRECISION_MAP[reader_precision]}`."
)
reader.to(PRECISION_MAP[reader_precision])
self.reader = reader
# self.reader = torch.compile(self.reader, backend="tvm")
# windowization stuff
self.tokenizer = SpacyTokenizer(language="en") # TODO: parametrize?
self.sentence_splitter: BaseSentenceSplitter | None = None
self.window_manager: WindowManager | None = None
if metadata_fields is None:
metadata_fields = []
self.metadata_fields = metadata_fields
# inference params
self.top_k = top_k
self.window_size = window_size
self.window_stride = window_stride
@staticmethod
def _instantiate_retriever(
retriever,
retriever_device,
retriever_precision,
document_index,
document_index_device,
document_index_precision,
):
if not isinstance(retriever, GoldenRetriever):
# convert to DictConfig
retriever = hydra.utils.instantiate(
OmegaConf.create(retriever),
device=retriever_device,
precision=retriever_precision,
index_device=document_index_device,
index_precision=document_index_precision,
)
retriever.training = False
retriever.eval()
if document_index is not None:
if retriever.document_index is not None:
logger.info(
"The Retriever already has a document index, replacing it with the provided one."
"If you want to keep using the old one, please do not provide a document index."
)
retriever.document_index = document_index
# we override the device and the precision of the document index if provided
if document_index_device is not None:
logger.info(f"Moving document index to `{document_index_device}`.")
retriever.document_index.to(document_index_device)
if document_index_precision is not None:
logger.info(
f"Setting precision of document index to `{PRECISION_MAP[document_index_precision]}`."
)
retriever.document_index.to(PRECISION_MAP[document_index_precision])
# retriever.document_index = document_index
# now we can move the retriever to the right device and set the precision
if retriever_device is not None:
logger.info(f"Moving retriever to `{retriever_device}`.")
retriever.to(retriever_device)
if retriever_precision is not None:
logger.info(
f"Setting precision of retriever to `{PRECISION_MAP[retriever_precision]}`."
)
retriever.to(PRECISION_MAP[retriever_precision])
return retriever
def __call__(
self,
text: str | List[str] | None = None,
windows: List[RelikReaderSample] | None = None,
candidates: List[str]
| List[Document]
| Dict[TaskType, List[Document]]
| None = None,
mentions: List[List[int]] | List[List[List[int]]] | None = None,
top_k: int | None = None,
window_size: int | None = None,
window_stride: int | None = None,
is_split_into_words: bool = False,
retriever_batch_size: int | None = 32,
reader_batch_size: int | None = 32,
return_also_windows: bool = False,
annotation_type: str | AnnotationType = AnnotationType.CHAR,
progress_bar: bool = False,
**kwargs,
) -> Union[RelikOutput, list[RelikOutput]]:
"""
Annotate a text with entities.
Args:
text (`str` or `list`):
The text to annotate. If a list is provided, each element of the list
will be annotated separately.
candidates (`list[str]`, `list[Document]`, `optional`, defaults to `None`):
The candidates to use for the reader. If `None`, the candidates will be
retrieved from the retriever.
mentions (`list[list[int]]` or `list[list[list[int]]]`, `optional`, defaults to `None`):
The mentions to use for the reader. If `None`, the mentions will be
predicted by the reader.
top_k (`int`, `optional`, defaults to `None`):
The number of candidates to retrieve for each window.
window_size (`int`, `optional`, defaults to `None`):
The size of the window. If `None`, the whole text will be annotated.
window_stride (`int`, `optional`, defaults to `None`):
The stride of the window. If `None`, there will be no overlap between windows.
retriever_batch_size (`int`, `optional`, defaults to `None`):
The batch size to use for the retriever. The whole input is the batch for the retriever.
reader_batch_size (`int`, `optional`, defaults to `None`):
The batch size to use for the reader. The whole input is the batch for the reader.
return_also_windows (`bool`, `optional`, defaults to `False`):
Whether to return the windows in the output.
annotation_type (`str` or `AnnotationType`, `optional`, defaults to `char`):
The type of annotation to return. If `char`, the spans will be in terms of
character offsets. If `word`, the spans will be in terms of word offsets.
**kwargs:
Additional keyword arguments to pass to the retriever and the reader.
Returns:
`RelikOutput` or `list[RelikOutput]`:
The annotated text. If a list was provided as input, a list of
`RelikOutput` objects will be returned.
"""
if text is None and windows is None:
raise ValueError(
"Either `text` or `windows` must be provided. Both are `None`."
)
if isinstance(annotation_type, str):
try:
annotation_type = AnnotationType(annotation_type)
except ValueError:
raise ValueError(
f"Annotation type {annotation_type} not recognized. "
f"Please choose one of {list(AnnotationType)}."
)
if top_k is None:
top_k = self.top_k or 100
if window_size is None:
window_size = self.window_size
if window_stride is None:
window_stride = self.window_stride
if text:
if isinstance(text, str):
text = [text]
if mentions is not None:
mentions = [mentions]
if file_logger is not None:
file_logger.info("Annotating the following text:")
for t in text:
file_logger.info(f" {t}")
if self.window_manager is None:
if window_size == "none":
self.sentence_splitter = BlankSentenceSplitter()
elif window_size == "sentence":
self.sentence_splitter = SpacySentenceSplitter()
else:
self.sentence_splitter = WindowSentenceSplitter(
window_size=window_size, window_stride=window_stride
)
self.window_manager = WindowManager(
self.tokenizer, self.sentence_splitter
)
if (
window_size not in ["sentence", "none"]
and window_stride is not None
and window_size < window_stride
):
raise ValueError(
f"Window size ({window_size}) must be greater than window stride ({window_stride})"
)
if windows is None:
# windows were provided, use them
windows, blank_windows = self.window_manager.create_windows(
text,
window_size,
window_stride,
is_split_into_words=is_split_into_words,
mentions=mentions
)
else:
blank_windows = []
text = {w.doc_id: w.text for w in windows}
if candidates is not None and any(
r is not None for r in self.retriever.values()
):
logger.info(
"Both candidates and a retriever were provided. "
"Retriever will be ignored."
)
windows_candidates = {TaskType.SPAN: None, TaskType.TRIPLET: None}
if candidates is not None:
# again, check if candidates is a dict
if isinstance(candidates, Dict):
if self.task not in candidates:
raise ValueError(
f"Task `{self.task}` not found in `candidates`."
f"Please choose one of {list(TaskType)}."
)
else:
candidates = {self.task: candidates}
for task_type, _candidates in candidates.items():
if isinstance(_candidates, list):
_candidates = [
[
c if isinstance(c, Document) else Document(c)
for c in _candidates[w.doc_id]
]
for w in windows
]
windows_candidates[task_type] = _candidates
else:
# retrieve candidates first
if self.retriever is None:
raise ValueError(
"No retriever was provided, please provide a retriever or candidates."
)
start_retr = time.time()
for task_type, retriever in self.retriever.items():
retriever_out = retriever.retrieve(
[w.text for w in windows],
text_pair=[w.doc_topic.text if w.doc_topic is not None else None for w in windows],
k=top_k,
batch_size=retriever_batch_size,
progress_bar=progress_bar,
**kwargs,
)
windows_candidates[task_type] = [
[p.document for p in predictions] for predictions in retriever_out
]
end_retr = time.time()
logger.info(f"Retrieval took {end_retr - start_retr} seconds.")
# clean up None's
windows_candidates = {
t: c for t, c in windows_candidates.items() if c is not None
}
# add passage to the windows
for task_type, task_candidates in windows_candidates.items():
for window, candidates in zip(windows, task_candidates):
# construct the candidates for the reader
formatted_candidates = []
for candidate in candidates:
window_candidate_text = candidate.text
for field in self.metadata_fields:
window_candidate_text += f"{candidate.metadata.get(field, '')}"
formatted_candidates.append(window_candidate_text)
# create a member for the windows that is named like the task
setattr(window, f"{task_type.value}_candidates", formatted_candidates)
for task_type, task_candidates in windows_candidates.items():
for window in blank_windows:
setattr(window, f"{task_type.value}_candidates", [])
setattr(window, "predicted_spans", [])
setattr(window, "predicted_triples", [])
if self.reader is not None:
start_read = time.time()
windows = self.reader.read(
samples=windows,
max_batch_size=reader_batch_size,
annotation_type=annotation_type,
progress_bar=progress_bar,
**kwargs,
)
end_read = time.time()
logger.info(f"Reading took {end_read - start_read} seconds.")
# TODO: check merging behavior without a reader
# do we want to merge windows if there is no reader?
if self.window_size is not None and self.window_size not in ["sentence", "none"]:
start_w = time.time()
windows = windows + blank_windows
windows.sort(key=lambda x: (x.doc_id, x.offset))
merged_windows = self.window_manager.merge_windows(windows)
end_w = time.time()
logger.info(f"Merging took {end_w - start_w} seconds.")
else:
merged_windows = windows
else:
windows = windows + blank_windows
windows.sort(key=lambda x: (x.doc_id, x.offset))
merged_windows = windows
# transform predictions into RelikOutput objects
output = []
for w in merged_windows:
span_labels = []
triples_labels = []
# span extraction should always be present
if getattr(w, "predicted_spans", None) is not None:
span_labels = sorted(
[
Span(start=ss, end=se, label=sl, text=text[w.doc_id][ss:se])
if annotation_type == AnnotationType.CHAR
else Span(start=ss, end=se, label=sl, text=w.words[ss:se])
for ss, se, sl in w.predicted_spans
],
key=lambda x: x.start,
)
# triple extraction is optional, if here add it
if getattr(w, "predicted_triples", None) is not None:
triples_labels = [
Triples(
subject=span_labels[subj],
label=label,
object=span_labels[obj],
confidence=conf,
)
for subj, label, obj, conf in w.predicted_triples
]
# create the output
sample_output = RelikOutput(
text=text[w.doc_id],
tokens=w.words,
spans=span_labels,
triples=triples_labels,
candidates={
task_type: [
r.document_index.documents.get_document_from_text(c)
for c in getattr(w, f"{task_type.value}_candidates", [])
if r.document_index.documents.get_document_from_text(c) is not None
]
for task_type, r in self.retriever.items()
},
)
output.append(sample_output)
# add windows to the output if requested
# do we want to force windows to be returned if there is no reader?
if return_also_windows:
for i, sample_output in enumerate(output):
sample_output.windows = [w for w in windows if w.doc_id == i]
# if only one text was provided, return a single RelikOutput object
if len(output) == 1:
return output[0]
return output
@classmethod
def from_pretrained(
cls,
model_name_or_dir: Union[str, os.PathLike],
config_file_name: str = CONFIG_NAME,
*args,
**kwargs,
) -> "Relik":
"""
Instantiate a `Relik` from a pretrained model.
Args:
model_name_or_dir (`str` or `os.PathLike`):
The name or path of the model to load.
config_file_name (`str`, `optional`, defaults to `config.yaml`):
The name of the configuration file to load.
*args:
Additional positional arguments to pass to `OmegaConf.merge`.
**kwargs:
Additional keyword arguments to pass to `OmegaConf.merge`.
Returns:
`Relik`:
The instantiated `Relik`.
"""
cache_dir = kwargs.pop("cache_dir", None)
force_download = kwargs.pop("force_download", False)
model_dir = from_cache(
model_name_or_dir,
filenames=[config_file_name],
cache_dir=cache_dir,
force_download=force_download,
)
config_path = model_dir / config_file_name
if not config_path.exists():
raise FileNotFoundError(
f"Model configuration file not found at {config_path}."
)
# overwrite config with config_kwargs
config = OmegaConf.load(config_path)
# if kwargs is not None:
config = OmegaConf.merge(config, OmegaConf.create(kwargs))
# do we want to print the config? I like it
logger.info(f"Loading Relik from {model_name_or_dir}")
# load relik from config
relik = hydra.utils.instantiate(config, _recursive_=False, *args)
return relik
def save_pretrained(
self,
output_dir: Union[str, os.PathLike],
config: Optional[Dict[str, Any]] = None,
config_file_name: Optional[str] = None,
save_weights: bool = False,
push_to_hub: bool = False,
model_id: Optional[str] = None,
organization: Optional[str] = None,
repo_name: Optional[str] = None,
retriever_model_id: Optional[str] = None,
reader_model_id: Optional[str] = None,
**kwargs,
):
"""
Save the configuration of Relik to the specified directory as a YAML file.
Args:
output_dir (`str`):
The directory to save the configuration file to.
config (`Optional[Dict[str, Any]]`, `optional`):
The configuration to save. If `None`, the current configuration will be
saved. Defaults to `None`.
config_file_name (`Optional[str]`, `optional`):
The name of the configuration file. Defaults to `config.yaml`.
save_weights (`bool`, `optional`):
Whether to save the weights of the model. Defaults to `False`.
push_to_hub (`bool`, `optional`):
Whether to push the saved model to the hub. Defaults to `False`.
model_id (`Optional[str]`, `optional`):
The id of the model to push to the hub. If `None`, the name of the
directory will be used. Defaults to `None`.
organization (`Optional[str]`, `optional`):
The organization to push the model to. Defaults to `None`.
repo_name (`Optional[str]`, `optional`):
The name of the repository to push the model to. Defaults to `None`.
retriever_model_id (`Optional[str]`, `optional`):
The id of the retriever model to push to the hub. If `None`, the name of the
directory will be used. Defaults to `None`.
reader_model_id (`Optional[str]`, `optional`):
The id of the reader model to push to the hub. If `None`, the name of the
directory will be used. Defaults to `None`.
**kwargs:
Additional keyword arguments to pass to `OmegaConf.save`.
"""
# create the output directory
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
retrievers_names: Dict[TaskType, Dict | None] = {
TaskType.SPAN: {
"question_encoder_name": None,
"passage_encoder_name": None,
"document_index_name": None,
},
TaskType.TRIPLET: {
"question_encoder_name": None,
"passage_encoder_name": None,
"document_index_name": None,
},
}
if save_weights:
# save weights
# retriever
model_id = model_id or output_dir.name
retriever_model_id = retriever_model_id or f"retriever-{model_id}"
for task_type, retriever in self.retriever.items():
if retriever is None:
continue
task_retriever_model_id = f"{retriever_model_id}-{task_type.value}"
question_encoder_name = f"{task_retriever_model_id}-question-encoder"
passage_encoder_name = f"{task_retriever_model_id}-passage-encoder"
document_index_name = f"{task_retriever_model_id}-index"
logger.info(
f"Saving retriever to {output_dir / task_retriever_model_id}"
)
retriever.save_pretrained(
output_dir / task_retriever_model_id,
question_encoder_name=question_encoder_name,
passage_encoder_name=passage_encoder_name,
document_index_name=document_index_name,
push_to_hub=push_to_hub,
organization=organization,
**kwargs,
)
retrievers_names[task_type] = {
"reader_model_id": task_retriever_model_id,
"question_encoder_name": question_encoder_name,
"passage_encoder_name": passage_encoder_name,
"document_index_name": document_index_name,
}
# reader
reader_model_id = reader_model_id or f"reader-{model_id}"
logger.info(f"Saving reader to {output_dir / reader_model_id}")
self.reader.save_pretrained(
output_dir / reader_model_id,
push_to_hub=push_to_hub,
organization=organization,
**kwargs,
)
if push_to_hub:
user = organization or get_logged_in_username()
# we need to update the config with the model ids that will
# result from the push to hub
for task_type, retriever_names in retrievers_names.items():
retriever_names[
"question_encoder_name"
] = f"{user}/{retriever_names['question_encoder_name']}"
retriever_names[
"passage_encoder_name"
] = f"{user}/{retriever_names['passage_encoder_name']}"
retriever_names[
"document_index_name"
] = f"{user}/{retriever_names['document_index_name']}"
# question_encoder_name = f"{user}/{question_encoder_name}"
# passage_encoder_name = f"{user}/{passage_encoder_name}"
# document_index_name = f"{user}/{document_index_name}"
reader_model_id = f"{user}/{reader_model_id}"
else:
for task_type, retriever_names in retrievers_names.items():
retriever_names["question_encoder_name"] = (
output_dir / retriever_names["question_encoder_name"]
)
retriever_names["passage_encoder_name"] = (
output_dir / retriever_names["passage_encoder_name"]
)
retriever_names["document_index_name"] = (
output_dir / retriever_names["document_index_name"]
)
reader_model_id = output_dir / reader_model_id
else:
# save config only
for task_type, retriever_names in retrievers_names.items():
retriever = self.retriever.get(task_type, None)
if retriever is None:
continue
retriever_names[
"question_encoder_name"
] = retriever.question_encoder.name_or_path
retriever_names[
"passage_encoder_name"
] = retriever.passage_encoder.name_or_path
retriever_names[
"document_index_name"
] = retriever.document_index.name_or_path
reader_model_id = self.reader.name_or_path
if config is None:
# create a default config
config = {
"_target_": f"{self.__class__.__module__}.{self.__class__.__name__}"
}
if self.retriever is not None:
config["retriever"] = {}
for task_type, retriever in self.retriever.items():
if retriever is None:
continue
config["retriever"][task_type.value] = {
"_target_": f"{retriever.__class__.__module__}.{retriever.__class__.__name__}",
}
if retriever.question_encoder is not None:
config["retriever"][task_type.value][
"question_encoder"
] = retrievers_names[task_type]["question_encoder_name"]
if (
retriever.passage_encoder is not None
and not retriever.passage_encoder_is_question_encoder
):
config["retriever"][task_type.value][
"passage_encoder"
] = retrievers_names[task_type]["passage_encoder_name"]
if retriever.document_index is not None:
config["retriever"][task_type.value][
"document_index"
] = retrievers_names[task_type]["document_index_name"]
if self.reader is not None:
config["reader"] = {
"_target_": f"{self.reader.__class__.__module__}.{self.reader.__class__.__name__}",
"transformer_model": reader_model_id,
}
# these are model-specific and should be saved
config["task"] = self.task
config["metadata_fields"] = self.metadata_fields
config["top_k"] = self.top_k
config["window_size"] = self.window_size
config["window_stride"] = self.window_stride
config_file_name = config_file_name or CONFIG_NAME
logger.info(f"Saving relik config to {output_dir / config_file_name}")
OmegaConf.save(config, output_dir / config_file_name)
if push_to_hub:
# push to hub
logger.info("Pushing to hub")
model_id = model_id or output_dir.name
upload(
output_dir,
model_id,
filenames=[config_file_name],
organization=organization,
repo_name=repo_name,
)