Update inference.py, for more general use.
Browse files- inference.py +2 -2
inference.py
CHANGED
@@ -61,8 +61,8 @@ class DualBertModel: # use this class to encode your document and queries.
|
|
61 |
with torch.no_grad():
|
62 |
to_return = []
|
63 |
for c in self.chunks(corpus, batch_size):
|
64 |
-
to_encode = [d['text'] for d in c]
|
65 |
-
ids, atm = self.process(
|
66 |
out = self.p_model(ids, atm)
|
67 |
to_return.append(out.cpu())
|
68 |
|
|
|
61 |
with torch.no_grad():
|
62 |
to_return = []
|
63 |
for c in self.chunks(corpus, batch_size):
|
64 |
+
# to_encode = [d['text'] for d in c]
|
65 |
+
ids, atm = self.process(c, False)
|
66 |
out = self.p_model(ids, atm)
|
67 |
to_return.append(out.cpu())
|
68 |
|