Remove Literal type hint (#156)

This commit is contained in:
Tanay Soni 2020-06-17 16:26:21 +02:00 committed by GitHub
parent af5fc79dce
commit 8bcc4b26a1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 3 additions and 4 deletions

View File

@ -45,7 +45,7 @@ if EMBEDDING_MODEL_PATH:
retriever = EmbeddingRetriever(
document_store=document_store,
embedding_model=EMBEDDING_MODEL_PATH,
model_format=EMBEDDING_MODEL_FORMAT, # type: ignore
model_format=EMBEDDING_MODEL_FORMAT,
gpu=USE_GPU
) # type: BaseRetriever
else:

View File

@ -2,7 +2,6 @@ import logging
from typing import List, Union
from farm.infer import Inferencer
from typing_extensions import Literal
from haystack.database.base import Document
from haystack.database.elasticsearch import ElasticsearchDocumentStore
@ -108,8 +107,8 @@ class EmbeddingRetriever(BaseRetriever):
document_store: ElasticsearchDocumentStore,
embedding_model: str,
gpu: bool = True,
model_format: Literal["farm", "transformers", "sentence_transformers"] = "farm",
pooling_strategy: Literal["cls_token", "reduce_mean", "reduce_max", "per_token", "s3e"] = "reduce_mean",
model_format: str = "farm",
pooling_strategy: str = "reduce_mean",
emb_extraction_layer: int = -1,
):
"""