Ramon Meffert
commited on
Commit
·
112c207
1
Parent(s):
ab5dfc2
Tweak reader parameters
Browse files
src/readers/dpr_reader.py
CHANGED
|
@@ -10,7 +10,10 @@ class DprReader():
|
|
| 10 |
"facebook/dpr-reader-single-nq-base"
|
| 11 |
)
|
| 12 |
|
| 13 |
-
def read(self,
|
|
|
|
|
|
|
|
|
|
| 14 |
encoded_inputs = self._tokenizer(
|
| 15 |
questions=query,
|
| 16 |
titles=context['titles'],
|
|
@@ -22,6 +25,10 @@ class DprReader():
|
|
| 22 |
outputs = self._model(**encoded_inputs)
|
| 23 |
|
| 24 |
predicted_spans = self._tokenizer.decode_best_spans(
|
| 25 |
-
encoded_inputs,
|
|
|
|
|
|
|
|
|
|
|
|
|
| 26 |
|
| 27 |
return predicted_spans
|
|
|
|
| 10 |
"facebook/dpr-reader-single-nq-base"
|
| 11 |
)
|
| 12 |
|
| 13 |
+
def read(self,
|
| 14 |
+
query: str,
|
| 15 |
+
context: Dict[str, List[str]],
|
| 16 |
+
num_answers=5) -> List[Tuple]:
|
| 17 |
encoded_inputs = self._tokenizer(
|
| 18 |
questions=query,
|
| 19 |
titles=context['titles'],
|
|
|
|
| 25 |
outputs = self._model(**encoded_inputs)
|
| 26 |
|
| 27 |
predicted_spans = self._tokenizer.decode_best_spans(
|
| 28 |
+
encoded_inputs,
|
| 29 |
+
outputs,
|
| 30 |
+
num_spans=num_answers,
|
| 31 |
+
num_spans_per_passage=2,
|
| 32 |
+
max_answer_length=512)
|
| 33 |
|
| 34 |
return predicted_spans
|