Skip to content

Change fixed value 256 in trange to batch_size #1401

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hanhainebula merged 9 commits into from
Apr 10, 2025
Merged

Change fixed value 256 in trange to batch_size #1401

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
005ffdf
change fixed disable limit to batch_size
pfldy2850 Mar 17, 2025
0f8b0d8
Update base.py
hanhainebula Apr 10, 2025
c43a3c6
Update base.py
hanhainebula Apr 10, 2025
dda2002
Update icl.py
hanhainebula Apr 10, 2025
df8b48c
Update m3.py
hanhainebula Apr 10, 2025
4e24ad4
Update base.py
hanhainebula Apr 10, 2025
cc978ae
Update layerwise.py
hanhainebula Apr 10, 2025
7a0a03a
Update lightweight.py
hanhainebula Apr 10, 2025
94068b2
Update base.py
hanhainebula Apr 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions FlagEmbedding/inference/embedder/decoder_only/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -225,7 +225,7 @@ def encode_single_device(
# tokenize without padding to get the correct length
all_inputs = []
for start_index in trange(0, len(sentences), batch_size, desc='pre tokenize',
disable=len(sentences) < 256):
disable=len(sentences) < batch_size):
sentences_batch = sentences[start_index:start_index + batch_size]
inputs_batch = self.tokenizer(
sentences_batch,
Expand Down Expand Up @@ -263,7 +263,7 @@ def encode_single_device(
# encode
all_embeddings = []
for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
disable=len(sentences) < 256):
disable=len(sentences) < batch_size):
inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
inputs_batch = self.tokenizer.pad(
inputs_batch,
Expand Down
9 changes: 5 additions & 4 deletions FlagEmbedding/inference/embedder/decoder_only/icl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -368,7 +368,8 @@ def encode_queries_single_device(

# tokenize without padding to get the correct length
all_inputs = []
for start_index in trange(0, len(input_texts), batch_size, desc='pre tokenize'):
for start_index in trange(0, len(input_texts), batch_size, desc='pre tokenize',
disable=len(input_texts) < batch_size):
sentences_batch = input_texts[start_index:start_index + batch_size]
inputs_batch = self.tokenizer(
sentences_batch,
Expand Down Expand Up @@ -417,7 +418,7 @@ def encode_queries_single_device(
# encode
all_embeddings = []
for start_index in tqdm(range(0, len(sentences_sorted), batch_size), desc="Inference Embeddings",
disable=len(sentences_sorted) < 256):
disable=len(sentences_sorted) < batch_size):
inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
inputs_batch = self.tokenizer.pad(
inputs_batch,
Expand Down Expand Up @@ -489,7 +490,7 @@ def encode_single_device(
# tokenize without padding to get the correct length
all_inputs = []
for start_index in trange(0, len(sentences), batch_size, desc='pre tokenize',
disable=len(sentences) < 256):
disable=len(sentences) < batch_size):
sentences_batch = sentences[start_index:start_index + batch_size]
inputs_batch = self.tokenizer(
sentences_batch,
Expand Down Expand Up @@ -527,7 +528,7 @@ def encode_single_device(
# encode
all_embeddings = []
for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
disable=len(sentences) < 256):
disable=len(sentences) < batch_size):
inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
inputs_batch = self.tokenizer.pad(
inputs_batch,
Expand Down
4 changes: 2 additions & 2 deletions FlagEmbedding/inference/embedder/encoder_only/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,7 +206,7 @@ def encode_single_device(
# tokenize without padding to get the correct length
all_inputs = []
for start_index in trange(0, len(sentences), batch_size, desc='pre tokenize',
disable=len(sentences) < 256):
disable=len(sentences) < batch_size):
sentences_batch = sentences[start_index:start_index + batch_size]
inputs_batch = self.tokenizer(
sentences_batch,
Expand Down Expand Up @@ -244,7 +244,7 @@ def encode_single_device(
# encode
all_embeddings = []
for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
disable=len(sentences) < 256):
disable=len(sentences) < batch_size):
inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
inputs_batch = self.tokenizer.pad(
inputs_batch,
Expand Down
6 changes: 3 additions & 3 deletions FlagEmbedding/inference/embedder/encoder_only/m3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -370,7 +370,7 @@ def _process_colbert_vecs(colbert_vecs: np.ndarray, attention_mask: list):
# tokenize without padding to get the correct length
all_inputs = []
for start_index in trange(0, len(sentences), batch_size, desc='pre tokenize',
disable=len(sentences) < 256):
disable=len(sentences) < batch_size):
sentences_batch = sentences[start_index:start_index + batch_size]
inputs_batch = self.tokenizer(
sentences_batch,
Expand Down Expand Up @@ -412,7 +412,7 @@ def _process_colbert_vecs(colbert_vecs: np.ndarray, attention_mask: list):
# encode
all_dense_embeddings, all_lexical_weights, all_colbert_vecs = [], [], []
for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
disable=len(sentences) < 256):
disable=len(sentences) < batch_size):
inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
inputs_batch = self.tokenizer.pad(
inputs_batch,
Expand Down Expand Up @@ -652,7 +652,7 @@ def _tokenize(texts: list, max_length: int):
'colbert+sparse+dense': []
}
for start_index in tqdm(range(0, len(sentence_pairs), batch_size), desc="Compute Scores",
disable=len(sentence_pairs) < 128):
disable=len(sentence_pairs) < batch_size):
sentences_batch = sentence_pairs[start_index:start_index + batch_size]

queries_batch = [pair[0] for pair in sentences_batch]
Expand Down
2 changes: 1 addition & 1 deletion FlagEmbedding/inference/reranker/decoder_only/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -310,7 +310,7 @@ def compute_score_single_gpu(
all_queries_inputs = []
all_passages_inputs = []
for start_index in trange(0, len(sentence_pairs), batch_size, desc="pre tokenize",
disable=len(sentence_pairs) < 128):
disable=len(sentence_pairs) < batch_size):
sentences_batch = sentence_pairs[start_index:start_index + batch_size]
queries = [s[0] for s in sentences_batch]
passages = [s[1] for s in sentences_batch]
Expand Down
2 changes: 1 addition & 1 deletion FlagEmbedding/inference/reranker/decoder_only/layerwise.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,7 +192,7 @@ def compute_score_single_gpu(
all_queries_inputs = []
all_passages_inputs = []
for start_index in trange(0, len(sentence_pairs), batch_size, desc="pre tokenize",
disable=len(sentence_pairs) < 128):
disable=len(sentence_pairs) < batch_size):
sentences_batch = sentence_pairs[start_index:start_index + batch_size]
queries = [s[0] for s in sentences_batch]
passages = [s[1] for s in sentences_batch]
Expand Down
2 changes: 1 addition & 1 deletion FlagEmbedding/inference/reranker/decoder_only/lightweight.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -263,7 +263,7 @@ def compute_score_single_gpu(
all_queries_inputs = []
all_passages_inputs = []
for start_index in trange(0, len(sentence_pairs), batch_size, desc="pre tokenize",
disable=len(sentence_pairs) < 128):
disable=len(sentence_pairs) < batch_size):
sentences_batch = sentence_pairs[start_index:start_index + batch_size]
queries = [s[0] for s in sentences_batch]
passages = [s[1] for s in sentences_batch]
Expand Down
4 changes: 2 additions & 2 deletions FlagEmbedding/inference/reranker/encoder_only/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ def compute_score_single_gpu(
# tokenize without padding to get the correct length
all_inputs = []
for start_index in trange(0, len(sentence_pairs), batch_size, desc="pre tokenize",
disable=len(sentence_pairs) < 128):
disable=len(sentence_pairs) < batch_size):
sentences_batch = sentence_pairs[start_index:start_index + batch_size]
queries = [s[0] for s in sentences_batch]
passages = [s[1] for s in sentences_batch]
Expand Down Expand Up @@ -174,7 +174,7 @@ def compute_score_single_gpu(

all_scores = []
for start_index in tqdm(range(0, len(all_inputs_sorted), batch_size), desc="Compute Scores",
disable=len(all_inputs_sorted) < 128):
disable=len(all_inputs_sorted) < batch_size):
sentences_batch = all_inputs_sorted[start_index:start_index + batch_size]
inputs = self.tokenizer.pad(
sentences_batch,
Expand Down