Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[RAG] Fix various generator issues #590

Merged
merged 3 commits into from
Nov 16, 2020
Merged

[RAG] Fix various generator issues #590

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
81c176e
Removing device information from generator model arguments as it is h…
lalitpagaria Nov 15, 2020
0c3f80c
num_return_sequences of should not be greate than num_beams
lalitpagaria Nov 15, 2020
bf18ba5
Raise error when user use generator with GPU as currently it is not s…
lalitpagaria Nov 16, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion haystack/generator/transformers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,6 @@ def __init__(
"""

self.model_name_or_path = model_name_or_path
self.top_k_answers = top_k_answers
self.max_length = max_length
self.min_length = min_length
self.generator_type = generator_type
Expand All @@ -99,8 +98,15 @@ def __init__(
self.prefix = prefix
self.retriever = retriever

if top_k_answers > self.num_beams:
top_k_answers = self.num_beams
logger.warning(f'top_k_answers value should not be greater than num_beams, hence setting it to {num_beams}')

self.top_k_answers = top_k_answers

if use_gpu and torch.cuda.is_available():
self.device = torch.device("cuda")
raise AttributeError("Currently RAGenerator does not support GPU, try with use_gpu=False")
else:
self.device = torch.device("cpu")

Expand Down Expand Up @@ -202,6 +208,11 @@ def predict(self, question: str, documents: List[Document], top_k: Optional[int]

top_k_answers = top_k if top_k is not None else self.top_k_answers

if top_k_answers > self.num_beams:
top_k_answers = self.num_beams
logger.warning(f'top_k_answers value should not be greater than num_beams, '
f'hence setting it to {top_k_answers}')

# Flatten the documents so easy to reference
flat_docs_dict: Dict[str, Any] = {}
for document in documents:
Expand Down