[RAG] Handle TFIDF retriever with pre-trained model (#4436)

* handle tfidf retriever from init pre-trained model

* needed this as well

* move change to TGA

Author: klshuster

parlai/agents/rag/retrievers.py

@@ -728,12 +728,17 @@ def __init__(self, opt: Opt, dictionary: DictionaryAgent, shared: TShared = None
         assert self.max_doc_paragraphs != 0
         if not shared:
             self.tfidf_retriever = create_agent(tfidf_opt)
+            self.query_encoder = DprQueryEncoder(
+                opt, dpr_model=opt['query_model'], pretrained_path=opt['dpr_model_file']
+            )
         else:
             self.tfidf_retriever = shared['tfidf_retriever']
+            self.query_encoder = shared['query_encoder']
 
     def share(self) -> TShared:
         shared = super().share()
         shared['tfidf_retriever'] = self.tfidf_retriever
+        shared['query_encoder'] = self.query_encoder
         return shared
 
     def retrieve_and_score(

parlai/core/torch_generator_agent.py

@@ -820,7 +820,7 @@ def _add_generation_metrics(self, batch, preds):
         """
         self.record_local_metric(
             'gen_n_toks',
-            AverageMetric.many([p.size(0) for p in preds], [1] * batch.batchsize),
+            AverageMetric.many([p.size(0) for p in preds], [1] * len(preds)),
         )
 
     def rank_eval_label_candidates(self, batch, batchsize):