from docproduct.predictor import GenerateQADoc

pretrained_path = 'pubmed_pmc_470k/'

ffn_weight_file = None

bert_ffn_weight_file = 'models/bertffn_crossentropy/bertffn'

embedding_file = 'qa_embeddings/bertffn_crossentropy.pkl'

doc = GenerateQADoc(pretrained_path=pretrained_path,

                    ffn_weight_file=None,

                    bert_ffn_weight_file=bert_ffn_weight_file,

                    embedding_file=embedding_file)

print(doc.predict('my eyes hurts and i have a headache.',

                  search_by='answer', topk=5, answer_only=False))

print(doc.predict('my eyes hurts and i have a headache.',

                  search_by='question', topk=5, answer_only=False))