Skip to content

Commit

Permalink
add rm tqdm (#47)
Browse files Browse the repository at this point in the history
adds progress bars to RM
  • Loading branch information
liana313 authored Dec 12, 2024
1 parent 733f253 commit b2246f5
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion lotus/models/sentence_transformers_rm.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import torch
from numpy.typing import NDArray
from sentence_transformers import SentenceTransformer
from tqdm import tqdm

from lotus.dtype_extensions import convert_to_base_data
from lotus.models.faiss_rm import FaissRM
Expand All @@ -27,7 +28,7 @@ def __init__(

def _embed(self, docs: pd.Series | list) -> NDArray[np.float64]:
all_embeddings = []
for i in range(0, len(docs), self.max_batch_size):
for i in tqdm(range(0, len(docs), self.max_batch_size)):
batch = docs[i : i + self.max_batch_size]
_batch = convert_to_base_data(batch)
torch_embeddings = self.transformer.encode(
Expand Down

0 comments on commit b2246f5

Please sign in to comment.