peter2000 commited on
Commit
fedecb4
1 Parent(s): 5c681e1

Update appStore/keyword_search.py

Browse files
Files changed (1) hide show
  1. appStore/keyword_search.py +2 -2
appStore/keyword_search.py CHANGED
@@ -15,7 +15,7 @@ import pandas as pd
15
  from sklearn.feature_extraction import _stop_words
16
  from haystack.document_stores import InMemoryDocumentStore
17
  from haystack.pipelines import ExtractiveQAPipeline
18
- from haystack.nodes import FARMReader, TfidfRetriever
19
  import string
20
  from markdown import markdown
21
  from annotated_text import annotation
@@ -31,7 +31,7 @@ logger = logging.getLogger(__name__)
31
  def start_haystack(documents_processed):
32
  document_store = InMemoryDocumentStore()
33
  document_store.write_documents(documents_processed)
34
- retriever = TfidfRetriever(document_store=document_store)
35
  reader = FARMReader(model_name_or_path="deepset/tinyroberta-squad2", use_gpu=True) #deepset/roberta-base-squad2
36
  pipeline = ExtractiveQAPipeline(reader, retriever)
37
  return pipeline
 
15
  from sklearn.feature_extraction import _stop_words
16
  from haystack.document_stores import InMemoryDocumentStore
17
  from haystack.pipelines import ExtractiveQAPipeline
18
+ from haystack.nodes import FARMReader, TfidfRetriever; EmbeddingRetriever
19
  import string
20
  from markdown import markdown
21
  from annotated_text import annotation
 
31
  def start_haystack(documents_processed):
32
  document_store = InMemoryDocumentStore()
33
  document_store.write_documents(documents_processed)
34
+ retriever = EmbeddingRetriever(document_store=document_store)
35
  reader = FARMReader(model_name_or_path="deepset/tinyroberta-squad2", use_gpu=True) #deepset/roberta-base-squad2
36
  pipeline = ExtractiveQAPipeline(reader, retriever)
37
  return pipeline