We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Wrong error message in tfidf retriever
https://github.com/facebookresearch/ParlAI/blob/24b49f1c6921cd2a80f4c99a980c438b318e2b5a/parlai/agents/tfidf_retriever/tfidf_retriever.py#L14

should check other places as well

(should be `pip install scikit-learn`)

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `parlai/agents/tfidf_retriever/tfidf_retriever.py`
Content:
```
1 #!/usr/bin/env python3
2 
3 # Copyright (c) 2017-present, Facebook, Inc.
4 # All rights reserved.
5 # This source code is licensed under the BSD-style license found in the
6 # LICENSE file in the root directory of this source tree. An additional grant
7 # of patent rights can be found in the PATENTS file in the same directory.
8 
9 try:
10     import regex  # noqa: F401
11     import scipy  # noqa: F401
12     import sklearn  # noqa: F401
13     import unicodedata  # noqa: F401
14     import pexpect  # noqa: F401
15 except ImportError:
16     raise ImportError('Please `pip install regex scipy sklearn pexpect`'
17                       ' to use the tfidf_retriever agent.')
18 
19 from parlai.core.agents import Agent
20 from parlai.core.utils import AttrDict
21 from .doc_db import DocDB
22 from .tfidf_doc_ranker import TfidfDocRanker
23 from .build_tfidf import run as build_tfidf
24 from .build_tfidf import live_count_matrix, get_tfidf_matrix
25 from numpy.random import choice
26 from collections import deque
27 import math
28 import random
29 import os
30 import pickle
31 import sqlite3
32 
33 
34 class TfidfRetrieverAgent(Agent):
35     """TFIDF-based retriever agent.
36 
37     If given a task to specify, will first store entries of that task into
38     a SQLite database and then build a sparse tfidf matrix of those entries.
39     If not, will build it on-the-fly whenever it sees observations with labels.
40 
41     This agent generates responses by building a sparse entry of the incoming
42     text observation, and then returning the highest-scoring documents
43     (calculated via sparse matrix multiplication) from the tfidf matrix.
44 
45     By default, this will return the "value" (the response) of the closest
46     entries. For example, saying "What is your favorite movie?" will not return
47     the text "Which movie is your favorite?" (high match) but rather the reply
48     to that (e.g. "Jurassic Park"). To use this agent for retrieving texts
49     (e.g. Wikipedia Entries), you can store label-less examples with the
50     '--retriever-task' argument and switch '--retriever-mode' to 'keys'.
51     """
52 
53     @staticmethod
54     def add_cmdline_args(parser):
55         parser = parser.add_argument_group('Retriever Arguments')
56         parser.add_argument(
57             '--retriever-numworkers', type=int, default=None,
58             help='Number of CPU processes (for tokenizing, etc)')
59         parser.add_argument(
60             '--retriever-ngram', type=int, default=2,
61             help='Use up to N-size n-grams (e.g. 2 = unigrams + bigrams)')
62         parser.add_argument(
63             '--retriever-hashsize', type=int, default=int(math.pow(2, 24)),
64             help='Number of buckets to use for hashing ngrams')
65         parser.add_argument(
66             '--retriever-tokenizer', type=str, default='simple',
67             help='String option specifying tokenizer type to use '
68                  '(e.g. "corenlp")')
69         parser.add_argument(
70             '--retriever-num-retrieved', default=5, type=int,
71             help='How many docs to retrieve.')
72         parser.add_argument(
73             '--retriever-mode', choices=['keys', 'values'], default='values',
74             help='Whether to retrieve the stored key or the stored value.'
75         )
76         parser.add_argument(
77             '--remove-title', type='bool', default=False,
78             help='Whether to remove the title from the retrieved passage')
79         parser.add_argument(
80             '--retriever-mode', choices=['keys', 'values'], default='values',
81             help='Whether to retrieve the stored key or the stored value. For '
82                  'example, if you want to return the text of an example, use '
83                  'keys here; if you want to return the label, use values here.'
84         )
85         parser.add_argument(
86             '--index-by-int-id', type='bool', default=True,
87             help='Whether to index into database by doc id as an integer. This \
88                   defaults to true for DBs built using ParlAI; for the DrQA \
89                   wiki dump, it is necessary to set this to False to \
90                   index into the DB appropriately')
91 
92     def __init__(self, opt, shared=None):
93         super().__init__(opt, shared)
94         self.id = 'SparseTfidfRetrieverAgent'
95         if not opt.get('model_file') or opt['model_file'] == '':
96             raise RuntimeError('Must set --model_file')
97 
98         opt['retriever_dbpath'] = opt['model_file'] + '.db'
99         opt['retriever_tfidfpath'] = opt['model_file'] + '.tfidf'
100 
101         self.db_path = opt['retriever_dbpath']
102         self.tfidf_path = opt['retriever_tfidfpath']
103 
104         self.tfidf_args = AttrDict({
105             'db_path': opt['retriever_dbpath'],
106             'out_dir': opt['retriever_tfidfpath'],
107             'ngram': opt['retriever_ngram'],
108             'hash_size': opt['retriever_hashsize'],
109             'tokenizer': opt['retriever_tokenizer'],
110             'num_workers': opt['retriever_numworkers'],
111         })
112 
113         if not os.path.exists(self.db_path):
114             conn = sqlite3.connect(self.db_path)
115             c = conn.cursor()
116             c.execute('CREATE TABLE documents '
117                       '(id INTEGER PRIMARY KEY, text, value);')
118             conn.commit()
119             conn.close()
120 
121         self.db = DocDB(db_path=opt['retriever_dbpath'])
122         if os.path.exists(self.tfidf_path + '.npz'):
123             self.ranker = TfidfDocRanker(
124                 tfidf_path=opt['retriever_tfidfpath'], strict=False)
125         self.ret_mode = opt['retriever_mode']
126         self.cands_hash = {}  # cache for candidates
127         self.triples_to_add = []  # in case we want to add more entries
128 
129         clen = opt.get('context_length', -1)
130         self.context_length = clen if clen >= 0 else None
131         self.include_labels = opt.get('include_labels', True)
132         self.reset()
133 
134     def reset(self):
135         super().reset()
136         self.episode_done = False
137         self.current = []
138         self.context = deque(maxlen=self.context_length)
139 
140     def train(self, mode=True):
141         self.training = mode
142 
143     def eval(self):
144         self.training = False
145 
146     def doc2txt(self, docid):
147         if not self.opt.get('index_by_int_id', True):
148             docid = self.ranker.get_doc_id(docid)
149         if self.ret_mode == 'keys':
150             return self.db.get_doc_text(docid)
151         elif self.ret_mode == 'values':
152             return self.db.get_doc_value(docid)
153         else:
154             raise RuntimeError('Retrieve mode {} not yet supported.'.format(
155                 self.ret_mode))
156 
157     def rebuild(self):
158         if len(self.triples_to_add) > 0:
159             self.db.add(self.triples_to_add)
160             self.triples_to_add.clear()
161             # rebuild tfidf
162             build_tfidf(self.tfidf_args)
163             self.ranker = TfidfDocRanker(
164                 tfidf_path=self.tfidf_path, strict=False)
165 
166     def save(self, path=None):
167         self.rebuild()
168         with open(self.opt['model_file'] + ".opt", 'wb') as handle:
169             pickle.dump(self.opt, handle, protocol=pickle.HIGHEST_PROTOCOL)
170         with open(self.opt['model_file'], 'w') as f:
171             f.write('\n')
172 
173     def train_act(self):
174         if ('ordered' not in self.opt.get('datatype', 'train:ordered') or
175                 self.opt.get('batchsize', 1) != 1 or
176                 self.opt.get('numthreads', 1) != 1 or
177                 self.opt.get('num_epochs', 1) != 1):
178             raise RuntimeError('Need to set --batchsize 1, --numthreads 1, \
179             --datatype train:ordered, --num_epochs 1')
180         obs = self.observation
181         self.current.append(obs)
182         self.episode_done = obs.get('episode_done', False)
183 
184         if self.episode_done:
185             for ex in self.current:
186                 if 'text' in ex:
187                     text = ex['text']
188                     self.context.append(text)
189                     if len(self.context) > 1:
190                         text = '\n'.join(self.context)
191 
192                 # add labels to context
193                 labels = ex.get('labels', ex.get('eval_labels'))
194                 label = None
195                 if labels is not None:
196                     label = random.choice(labels)
197                     if self.include_labels:
198                         self.context.append(label)
199                 # use None for ID to auto-assign doc ids--we don't need to
200                 # ever reverse-lookup them
201                 self.triples_to_add.append((None, text, label))
202 
203             self.episode_done = False
204             self.current.clear()
205             self.context.clear()
206 
207         return {'id': self.getID(),
208                 'text': obs.get('labels', ['I don\'t know'])[0]}
209 
210     def act(self):
211         obs = self.observation
212         reply = {}
213         reply['id'] = self.getID()
214         if 'labels' in obs:
215             return self.train_act()
216         if 'text' in obs:
217             self.rebuild()  # no-op if nothing has been queued to store
218             doc_ids, doc_scores = self.ranker.closest_docs(
219                 obs['text'],
220                 self.opt.get('retriever_num_retrieved', 5)
221             )
222 
223             if False and obs.get('label_candidates'):  # TODO: Alex (doesn't work)
224                 # these are better selection than stored facts
225                 # rank these options instead
226                 cands = obs['label_candidates']
227                 cands_id = id(cands)
228                 if cands_id not in self.cands_hash:
229                     # cache candidate set
230                     # will not update if cand set changes contents
231                     c_list = list(cands)
232                     self.cands_hash[cands_id] = (
233                         get_tfidf_matrix(
234                             live_count_matrix(self.tfidf_args, c_list)
235                         ),
236                         c_list
237                     )
238                 c_ids, c_scores = self.ranker.closest_docs(
239                     obs['text'],
240                     self.opt.get('retriever_num_retrieved', 5),
241                     matrix=self.cands_hash[cands_id][0]
242                 )
243                 reply['text_candidates'] = [
244                     self.cands_hash[cands_id][1][cid] for cid in c_ids
245                 ]
246                 reply['candidate_scores'] = c_scores
247                 if len(reply['text_candidates']) > 0:
248                     reply['text'] = reply['text_candidates'][0]
249                 else:
250                     reply['text'] = ''
251             elif len(doc_ids) > 0:
252                 # return stored fact
253                 # total = sum(doc_scores)
254                 # doc_probs = [d / total for d in doc_scores]
255 
256                 # returned
257                 picks = [self.doc2txt(int(did)) for did in doc_ids]
258                 pick = self.doc2txt(int(doc_ids[0]))  # select best response
259 
260                 if self.opt.get('remove_title', False):
261                     picks = ['\n'.join(p.split('\n')[1:]) for p in picks]
262                     pick = '\n'.join(pick.split('\n')[1:])
263                 reply['text_candidates'] = picks
264                 reply['candidate_scores'] = doc_scores
265 
266                 # could pick single choice based on probability scores?
267                 # pick = int(choice(doc_ids, p=doc_probs))
268                 reply['text'] = pick
269             else:
270                 # no cands and nothing found, return generic response
271                 reply['text'] = choice([
272                     'Can you say something more interesting?',
273                     'Why are you being so short with me?',
274                     'What are you really thinking?',
275                     'Can you expand on that?',
276                 ])
277 
278         return reply
279
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/parlai/agents/tfidf_retriever/tfidf_retriever.py b/parlai/agents/tfidf_retriever/tfidf_retriever.py
--- a/parlai/agents/tfidf_retriever/tfidf_retriever.py
+++ b/parlai/agents/tfidf_retriever/tfidf_retriever.py
@@ -13,7 +13,7 @@
     import unicodedata  # noqa: F401
     import pexpect  # noqa: F401
 except ImportError:
-    raise ImportError('Please `pip install regex scipy sklearn pexpect`'
+    raise ImportError('Please `pip install regex scipy scikit-learn pexpect`'
                       ' to use the tfidf_retriever agent.')
 
 from parlai.core.agents import Agent

{"golden_diff": "diff --git a/parlai/agents/tfidf_retriever/tfidf_retriever.py b/parlai/agents/tfidf_retriever/tfidf_retriever.py\n--- a/parlai/agents/tfidf_retriever/tfidf_retriever.py\n+++ b/parlai/agents/tfidf_retriever/tfidf_retriever.py\n@@ -13,7 +13,7 @@\n     import unicodedata  # noqa: F401\n     import pexpect  # noqa: F401\n except ImportError:\n-    raise ImportError('Please `pip install regex scipy sklearn pexpect`'\n+    raise ImportError('Please `pip install regex scipy scikit-learn pexpect`'\n                       ' to use the tfidf_retriever agent.')\n \n from parlai.core.agents import Agent\n", "issue": "Wrong error message in tfidf retriever\nhttps://github.com/facebookresearch/ParlAI/blob/24b49f1c6921cd2a80f4c99a980c438b318e2b5a/parlai/agents/tfidf_retriever/tfidf_retriever.py#L14\r\n\r\nshould check other places as well\r\n\r\n(should be `pip install scikit-learn`)\n", "before_files": [{"content": "#!/usr/bin/env python3\n\n# Copyright (c) 2017-present, Facebook, Inc.\n# All rights reserved.\n# This source code is licensed under the BSD-style license found in the\n# LICENSE file in the root directory of this source tree. An additional grant\n# of patent rights can be found in the PATENTS file in the same directory.\n\ntry:\n    import regex  # noqa: F401\n    import scipy  # noqa: F401\n    import sklearn  # noqa: F401\n    import unicodedata  # noqa: F401\n    import pexpect  # noqa: F401\nexcept ImportError:\n    raise ImportError('Please `pip install regex scipy sklearn pexpect`'\n                      ' to use the tfidf_retriever agent.')\n\nfrom parlai.core.agents import Agent\nfrom parlai.core.utils import AttrDict\nfrom .doc_db import DocDB\nfrom .tfidf_doc_ranker import TfidfDocRanker\nfrom .build_tfidf import run as build_tfidf\nfrom .build_tfidf import live_count_matrix, get_tfidf_matrix\nfrom numpy.random import choice\nfrom collections import deque\nimport math\nimport random\nimport os\nimport pickle\nimport sqlite3\n\n\nclass TfidfRetrieverAgent(Agent):\n    \"\"\"TFIDF-based retriever agent.\n\n    If given a task to specify, will first store entries of that task into\n    a SQLite database and then build a sparse tfidf matrix of those entries.\n    If not, will build it on-the-fly whenever it sees observations with labels.\n\n    This agent generates responses by building a sparse entry of the incoming\n    text observation, and then returning the highest-scoring documents\n    (calculated via sparse matrix multiplication) from the tfidf matrix.\n\n    By default, this will return the \"value\" (the response) of the closest\n    entries. For example, saying \"What is your favorite movie?\" will not return\n    the text \"Which movie is your favorite?\" (high match) but rather the reply\n    to that (e.g. \"Jurassic Park\"). To use this agent for retrieving texts\n    (e.g. Wikipedia Entries), you can store label-less examples with the\n    '--retriever-task' argument and switch '--retriever-mode' to 'keys'.\n    \"\"\"\n\n    @staticmethod\n    def add_cmdline_args(parser):\n        parser = parser.add_argument_group('Retriever Arguments')\n        parser.add_argument(\n            '--retriever-numworkers', type=int, default=None,\n            help='Number of CPU processes (for tokenizing, etc)')\n        parser.add_argument(\n            '--retriever-ngram', type=int, default=2,\n            help='Use up to N-size n-grams (e.g. 2 = unigrams + bigrams)')\n        parser.add_argument(\n            '--retriever-hashsize', type=int, default=int(math.pow(2, 24)),\n            help='Number of buckets to use for hashing ngrams')\n        parser.add_argument(\n            '--retriever-tokenizer', type=str, default='simple',\n            help='String option specifying tokenizer type to use '\n                 '(e.g. \"corenlp\")')\n        parser.add_argument(\n            '--retriever-num-retrieved', default=5, type=int,\n            help='How many docs to retrieve.')\n        parser.add_argument(\n            '--retriever-mode', choices=['keys', 'values'], default='values',\n            help='Whether to retrieve the stored key or the stored value.'\n        )\n        parser.add_argument(\n            '--remove-title', type='bool', default=False,\n            help='Whether to remove the title from the retrieved passage')\n        parser.add_argument(\n            '--retriever-mode', choices=['keys', 'values'], default='values',\n            help='Whether to retrieve the stored key or the stored value. For '\n                 'example, if you want to return the text of an example, use '\n                 'keys here; if you want to return the label, use values here.'\n        )\n        parser.add_argument(\n            '--index-by-int-id', type='bool', default=True,\n            help='Whether to index into database by doc id as an integer. This \\\n                  defaults to true for DBs built using ParlAI; for the DrQA \\\n                  wiki dump, it is necessary to set this to False to \\\n                  index into the DB appropriately')\n\n    def __init__(self, opt, shared=None):\n        super().__init__(opt, shared)\n        self.id = 'SparseTfidfRetrieverAgent'\n        if not opt.get('model_file') or opt['model_file'] == '':\n            raise RuntimeError('Must set --model_file')\n\n        opt['retriever_dbpath'] = opt['model_file'] + '.db'\n        opt['retriever_tfidfpath'] = opt['model_file'] + '.tfidf'\n\n        self.db_path = opt['retriever_dbpath']\n        self.tfidf_path = opt['retriever_tfidfpath']\n\n        self.tfidf_args = AttrDict({\n            'db_path': opt['retriever_dbpath'],\n            'out_dir': opt['retriever_tfidfpath'],\n            'ngram': opt['retriever_ngram'],\n            'hash_size': opt['retriever_hashsize'],\n            'tokenizer': opt['retriever_tokenizer'],\n            'num_workers': opt['retriever_numworkers'],\n        })\n\n        if not os.path.exists(self.db_path):\n            conn = sqlite3.connect(self.db_path)\n            c = conn.cursor()\n            c.execute('CREATE TABLE documents '\n                      '(id INTEGER PRIMARY KEY, text, value);')\n            conn.commit()\n            conn.close()\n\n        self.db = DocDB(db_path=opt['retriever_dbpath'])\n        if os.path.exists(self.tfidf_path + '.npz'):\n            self.ranker = TfidfDocRanker(\n                tfidf_path=opt['retriever_tfidfpath'], strict=False)\n        self.ret_mode = opt['retriever_mode']\n        self.cands_hash = {}  # cache for candidates\n        self.triples_to_add = []  # in case we want to add more entries\n\n        clen = opt.get('context_length', -1)\n        self.context_length = clen if clen >= 0 else None\n        self.include_labels = opt.get('include_labels', True)\n        self.reset()\n\n    def reset(self):\n        super().reset()\n        self.episode_done = False\n        self.current = []\n        self.context = deque(maxlen=self.context_length)\n\n    def train(self, mode=True):\n        self.training = mode\n\n    def eval(self):\n        self.training = False\n\n    def doc2txt(self, docid):\n        if not self.opt.get('index_by_int_id', True):\n            docid = self.ranker.get_doc_id(docid)\n        if self.ret_mode == 'keys':\n            return self.db.get_doc_text(docid)\n        elif self.ret_mode == 'values':\n            return self.db.get_doc_value(docid)\n        else:\n            raise RuntimeError('Retrieve mode {} not yet supported.'.format(\n                self.ret_mode))\n\n    def rebuild(self):\n        if len(self.triples_to_add) > 0:\n            self.db.add(self.triples_to_add)\n            self.triples_to_add.clear()\n            # rebuild tfidf\n            build_tfidf(self.tfidf_args)\n            self.ranker = TfidfDocRanker(\n                tfidf_path=self.tfidf_path, strict=False)\n\n    def save(self, path=None):\n        self.rebuild()\n        with open(self.opt['model_file'] + \".opt\", 'wb') as handle:\n            pickle.dump(self.opt, handle, protocol=pickle.HIGHEST_PROTOCOL)\n        with open(self.opt['model_file'], 'w') as f:\n            f.write('\\n')\n\n    def train_act(self):\n        if ('ordered' not in self.opt.get('datatype', 'train:ordered') or\n                self.opt.get('batchsize', 1) != 1 or\n                self.opt.get('numthreads', 1) != 1 or\n                self.opt.get('num_epochs', 1) != 1):\n            raise RuntimeError('Need to set --batchsize 1, --numthreads 1, \\\n            --datatype train:ordered, --num_epochs 1')\n        obs = self.observation\n        self.current.append(obs)\n        self.episode_done = obs.get('episode_done', False)\n\n        if self.episode_done:\n            for ex in self.current:\n                if 'text' in ex:\n                    text = ex['text']\n                    self.context.append(text)\n                    if len(self.context) > 1:\n                        text = '\\n'.join(self.context)\n\n                # add labels to context\n                labels = ex.get('labels', ex.get('eval_labels'))\n                label = None\n                if labels is not None:\n                    label = random.choice(labels)\n                    if self.include_labels:\n                        self.context.append(label)\n                # use None for ID to auto-assign doc ids--we don't need to\n                # ever reverse-lookup them\n                self.triples_to_add.append((None, text, label))\n\n            self.episode_done = False\n            self.current.clear()\n            self.context.clear()\n\n        return {'id': self.getID(),\n                'text': obs.get('labels', ['I don\\'t know'])[0]}\n\n    def act(self):\n        obs = self.observation\n        reply = {}\n        reply['id'] = self.getID()\n        if 'labels' in obs:\n            return self.train_act()\n        if 'text' in obs:\n            self.rebuild()  # no-op if nothing has been queued to store\n            doc_ids, doc_scores = self.ranker.closest_docs(\n                obs['text'],\n                self.opt.get('retriever_num_retrieved', 5)\n            )\n\n            if False and obs.get('label_candidates'):  # TODO: Alex (doesn't work)\n                # these are better selection than stored facts\n                # rank these options instead\n                cands = obs['label_candidates']\n                cands_id = id(cands)\n                if cands_id not in self.cands_hash:\n                    # cache candidate set\n                    # will not update if cand set changes contents\n                    c_list = list(cands)\n                    self.cands_hash[cands_id] = (\n                        get_tfidf_matrix(\n                            live_count_matrix(self.tfidf_args, c_list)\n                        ),\n                        c_list\n                    )\n                c_ids, c_scores = self.ranker.closest_docs(\n                    obs['text'],\n                    self.opt.get('retriever_num_retrieved', 5),\n                    matrix=self.cands_hash[cands_id][0]\n                )\n                reply['text_candidates'] = [\n                    self.cands_hash[cands_id][1][cid] for cid in c_ids\n                ]\n                reply['candidate_scores'] = c_scores\n                if len(reply['text_candidates']) > 0:\n                    reply['text'] = reply['text_candidates'][0]\n                else:\n                    reply['text'] = ''\n            elif len(doc_ids) > 0:\n                # return stored fact\n                # total = sum(doc_scores)\n                # doc_probs = [d / total for d in doc_scores]\n\n                # returned\n                picks = [self.doc2txt(int(did)) for did in doc_ids]\n                pick = self.doc2txt(int(doc_ids[0]))  # select best response\n\n                if self.opt.get('remove_title', False):\n                    picks = ['\\n'.join(p.split('\\n')[1:]) for p in picks]\n                    pick = '\\n'.join(pick.split('\\n')[1:])\n                reply['text_candidates'] = picks\n                reply['candidate_scores'] = doc_scores\n\n                # could pick single choice based on probability scores?\n                # pick = int(choice(doc_ids, p=doc_probs))\n                reply['text'] = pick\n            else:\n                # no cands and nothing found, return generic response\n                reply['text'] = choice([\n                    'Can you say something more interesting?',\n                    'Why are you being so short with me?',\n                    'What are you really thinking?',\n                    'Can you expand on that?',\n                ])\n\n        return reply\n", "path": "parlai/agents/tfidf_retriever/tfidf_retriever.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n\n# Copyright (c) 2017-present, Facebook, Inc.\n# All rights reserved.\n# This source code is licensed under the BSD-style license found in the\n# LICENSE file in the root directory of this source tree. An additional grant\n# of patent rights can be found in the PATENTS file in the same directory.\n\ntry:\n    import regex  # noqa: F401\n    import scipy  # noqa: F401\n    import sklearn  # noqa: F401\n    import unicodedata  # noqa: F401\n    import pexpect  # noqa: F401\nexcept ImportError:\n    raise ImportError('Please `pip install regex scipy scikit-learn pexpect`'\n                      ' to use the tfidf_retriever agent.')\n\nfrom parlai.core.agents import Agent\nfrom parlai.core.utils import AttrDict\nfrom .doc_db import DocDB\nfrom .tfidf_doc_ranker import TfidfDocRanker\nfrom .build_tfidf import run as build_tfidf\nfrom .build_tfidf import live_count_matrix, get_tfidf_matrix\nfrom numpy.random import choice\nfrom collections import deque\nimport math\nimport random\nimport os\nimport pickle\nimport sqlite3\n\n\nclass TfidfRetrieverAgent(Agent):\n    \"\"\"TFIDF-based retriever agent.\n\n    If given a task to specify, will first store entries of that task into\n    a SQLite database and then build a sparse tfidf matrix of those entries.\n    If not, will build it on-the-fly whenever it sees observations with labels.\n\n    This agent generates responses by building a sparse entry of the incoming\n    text observation, and then returning the highest-scoring documents\n    (calculated via sparse matrix multiplication) from the tfidf matrix.\n\n    By default, this will return the \"value\" (the response) of the closest\n    entries. For example, saying \"What is your favorite movie?\" will not return\n    the text \"Which movie is your favorite?\" (high match) but rather the reply\n    to that (e.g. \"Jurassic Park\"). To use this agent for retrieving texts\n    (e.g. Wikipedia Entries), you can store label-less examples with the\n    '--retriever-task' argument and switch '--retriever-mode' to 'keys'.\n    \"\"\"\n\n    @staticmethod\n    def add_cmdline_args(parser):\n        parser = parser.add_argument_group('Retriever Arguments')\n        parser.add_argument(\n            '--retriever-numworkers', type=int, default=None,\n            help='Number of CPU processes (for tokenizing, etc)')\n        parser.add_argument(\n            '--retriever-ngram', type=int, default=2,\n            help='Use up to N-size n-grams (e.g. 2 = unigrams + bigrams)')\n        parser.add_argument(\n            '--retriever-hashsize', type=int, default=int(math.pow(2, 24)),\n            help='Number of buckets to use for hashing ngrams')\n        parser.add_argument(\n            '--retriever-tokenizer', type=str, default='simple',\n            help='String option specifying tokenizer type to use '\n                 '(e.g. \"corenlp\")')\n        parser.add_argument(\n            '--retriever-num-retrieved', default=5, type=int,\n            help='How many docs to retrieve.')\n        parser.add_argument(\n            '--retriever-mode', choices=['keys', 'values'], default='values',\n            help='Whether to retrieve the stored key or the stored value.'\n        )\n        parser.add_argument(\n            '--remove-title', type='bool', default=False,\n            help='Whether to remove the title from the retrieved passage')\n        parser.add_argument(\n            '--retriever-mode', choices=['keys', 'values'], default='values',\n            help='Whether to retrieve the stored key or the stored value. For '\n                 'example, if you want to return the text of an example, use '\n                 'keys here; if you want to return the label, use values here.'\n        )\n        parser.add_argument(\n            '--index-by-int-id', type='bool', default=True,\n            help='Whether to index into database by doc id as an integer. This \\\n                  defaults to true for DBs built using ParlAI; for the DrQA \\\n                  wiki dump, it is necessary to set this to False to \\\n                  index into the DB appropriately')\n\n    def __init__(self, opt, shared=None):\n        super().__init__(opt, shared)\n        self.id = 'SparseTfidfRetrieverAgent'\n        if not opt.get('model_file') or opt['model_file'] == '':\n            raise RuntimeError('Must set --model_file')\n\n        opt['retriever_dbpath'] = opt['model_file'] + '.db'\n        opt['retriever_tfidfpath'] = opt['model_file'] + '.tfidf'\n\n        self.db_path = opt['retriever_dbpath']\n        self.tfidf_path = opt['retriever_tfidfpath']\n\n        self.tfidf_args = AttrDict({\n            'db_path': opt['retriever_dbpath'],\n            'out_dir': opt['retriever_tfidfpath'],\n            'ngram': opt['retriever_ngram'],\n            'hash_size': opt['retriever_hashsize'],\n            'tokenizer': opt['retriever_tokenizer'],\n            'num_workers': opt['retriever_numworkers'],\n        })\n\n        if not os.path.exists(self.db_path):\n            conn = sqlite3.connect(self.db_path)\n            c = conn.cursor()\n            c.execute('CREATE TABLE documents '\n                      '(id INTEGER PRIMARY KEY, text, value);')\n            conn.commit()\n            conn.close()\n\n        self.db = DocDB(db_path=opt['retriever_dbpath'])\n        if os.path.exists(self.tfidf_path + '.npz'):\n            self.ranker = TfidfDocRanker(\n                tfidf_path=opt['retriever_tfidfpath'], strict=False)\n        self.ret_mode = opt['retriever_mode']\n        self.cands_hash = {}  # cache for candidates\n        self.triples_to_add = []  # in case we want to add more entries\n\n        clen = opt.get('context_length', -1)\n        self.context_length = clen if clen >= 0 else None\n        self.include_labels = opt.get('include_labels', True)\n        self.reset()\n\n    def reset(self):\n        super().reset()\n        self.episode_done = False\n        self.current = []\n        self.context = deque(maxlen=self.context_length)\n\n    def train(self, mode=True):\n        self.training = mode\n\n    def eval(self):\n        self.training = False\n\n    def doc2txt(self, docid):\n        if not self.opt.get('index_by_int_id', True):\n            docid = self.ranker.get_doc_id(docid)\n        if self.ret_mode == 'keys':\n            return self.db.get_doc_text(docid)\n        elif self.ret_mode == 'values':\n            return self.db.get_doc_value(docid)\n        else:\n            raise RuntimeError('Retrieve mode {} not yet supported.'.format(\n                self.ret_mode))\n\n    def rebuild(self):\n        if len(self.triples_to_add) > 0:\n            self.db.add(self.triples_to_add)\n            self.triples_to_add.clear()\n            # rebuild tfidf\n            build_tfidf(self.tfidf_args)\n            self.ranker = TfidfDocRanker(\n                tfidf_path=self.tfidf_path, strict=False)\n\n    def save(self, path=None):\n        self.rebuild()\n        with open(self.opt['model_file'] + \".opt\", 'wb') as handle:\n            pickle.dump(self.opt, handle, protocol=pickle.HIGHEST_PROTOCOL)\n        with open(self.opt['model_file'], 'w') as f:\n            f.write('\\n')\n\n    def train_act(self):\n        if ('ordered' not in self.opt.get('datatype', 'train:ordered') or\n                self.opt.get('batchsize', 1) != 1 or\n                self.opt.get('numthreads', 1) != 1 or\n                self.opt.get('num_epochs', 1) != 1):\n            raise RuntimeError('Need to set --batchsize 1, --numthreads 1, \\\n            --datatype train:ordered, --num_epochs 1')\n        obs = self.observation\n        self.current.append(obs)\n        self.episode_done = obs.get('episode_done', False)\n\n        if self.episode_done:\n            for ex in self.current:\n                if 'text' in ex:\n                    text = ex['text']\n                    self.context.append(text)\n                    if len(self.context) > 1:\n                        text = '\\n'.join(self.context)\n\n                # add labels to context\n                labels = ex.get('labels', ex.get('eval_labels'))\n                label = None\n                if labels is not None:\n                    label = random.choice(labels)\n                    if self.include_labels:\n                        self.context.append(label)\n                # use None for ID to auto-assign doc ids--we don't need to\n                # ever reverse-lookup them\n                self.triples_to_add.append((None, text, label))\n\n            self.episode_done = False\n            self.current.clear()\n            self.context.clear()\n\n        return {'id': self.getID(),\n                'text': obs.get('labels', ['I don\\'t know'])[0]}\n\n    def act(self):\n        obs = self.observation\n        reply = {}\n        reply['id'] = self.getID()\n        if 'labels' in obs:\n            return self.train_act()\n        if 'text' in obs:\n            self.rebuild()  # no-op if nothing has been queued to store\n            doc_ids, doc_scores = self.ranker.closest_docs(\n                obs['text'],\n                self.opt.get('retriever_num_retrieved', 5)\n            )\n\n            if False and obs.get('label_candidates'):  # TODO: Alex (doesn't work)\n                # these are better selection than stored facts\n                # rank these options instead\n                cands = obs['label_candidates']\n                cands_id = id(cands)\n                if cands_id not in self.cands_hash:\n                    # cache candidate set\n                    # will not update if cand set changes contents\n                    c_list = list(cands)\n                    self.cands_hash[cands_id] = (\n                        get_tfidf_matrix(\n                            live_count_matrix(self.tfidf_args, c_list)\n                        ),\n                        c_list\n                    )\n                c_ids, c_scores = self.ranker.closest_docs(\n                    obs['text'],\n                    self.opt.get('retriever_num_retrieved', 5),\n                    matrix=self.cands_hash[cands_id][0]\n                )\n                reply['text_candidates'] = [\n                    self.cands_hash[cands_id][1][cid] for cid in c_ids\n                ]\n                reply['candidate_scores'] = c_scores\n                if len(reply['text_candidates']) > 0:\n                    reply['text'] = reply['text_candidates'][0]\n                else:\n                    reply['text'] = ''\n            elif len(doc_ids) > 0:\n                # return stored fact\n                # total = sum(doc_scores)\n                # doc_probs = [d / total for d in doc_scores]\n\n                # returned\n                picks = [self.doc2txt(int(did)) for did in doc_ids]\n                pick = self.doc2txt(int(doc_ids[0]))  # select best response\n\n                if self.opt.get('remove_title', False):\n                    picks = ['\\n'.join(p.split('\\n')[1:]) for p in picks]\n                    pick = '\\n'.join(pick.split('\\n')[1:])\n                reply['text_candidates'] = picks\n                reply['candidate_scores'] = doc_scores\n\n                # could pick single choice based on probability scores?\n                # pick = int(choice(doc_ids, p=doc_probs))\n                reply['text'] = pick\n            else:\n                # no cands and nothing found, return generic response\n                reply['text'] = choice([\n                    'Can you say something more interesting?',\n                    'Why are you being so short with me?',\n                    'What are you really thinking?',\n                    'Can you expand on that?',\n                ])\n\n        return reply\n", "path": "parlai/agents/tfidf_retriever/tfidf_retriever.py"}]}