Merge branch 'main' into Feature/#335

autorag/nodes/passagereranker/monot5.py

@@ -1,10 +1,12 @@
-import asyncio
+from itertools import chain
 from typing import List, Tuple
 
+import pandas as pd
 import torch
 from transformers import T5Tokenizer, T5ForConditionalGeneration
 
 from autorag.nodes.passagereranker.base import passage_reranker_node
+from autorag.utils.util import make_batch, sort_by_scores, flatten_apply, select_top_k
 
 prediction_tokens = {
     'castorini/monot5-base-msmarco': ['▁false', '▁true'],
@@ -33,7 +35,8 @@
 @passage_reranker_node
 def monot5(queries: List[str], contents_list: List[List[str]],
            scores_list: List[List[float]], ids_list: List[List[str]],
-           top_k: int, model_name: str = 'castorini/monot5-3b-msmarco-10k') \
+           top_k: int, model_name: str = 'castorini/monot5-3b-msmarco-10k',
+           batch: int = 64, ) \
         -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
     """
     Rerank a list of contents based on their relevance to a query using MonoT5.
@@ -48,6 +51,7 @@ def monot5(queries: List[str], contents_list: List[List[str]],
             If there is a '/' in the model name parameter,
             when we create the file to store the results, the path will be twisted because of the '/'.
             Therefore, it will be received as '_' instead of '/'.
+    :param batch: The number of queries to be processed in a batch
     :return: tuple of lists containing the reranked contents, ids, and scores
     """
     # replace '_' to '/'
@@ -61,72 +65,47 @@ def monot5(queries: List[str], contents_list: List[List[str]],
     token_false_id = tokenizer.convert_tokens_to_ids(token_false)
     token_true_id = tokenizer.convert_tokens_to_ids(token_true)
     # Determine the device to run the model on (GPU if available, otherwise CPU)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     model.to(device)
-    # Run async mono_t5_rerank_pure function
-    tasks = [mono_t5_pure(query, contents, scores, top_k, ids, model, device, tokenizer, token_false_id, token_true_id) \
-             for query, contents, scores, ids in zip(queries, contents_list, scores_list, ids_list)]
-    loop = asyncio.get_event_loop()
-    results = loop.run_until_complete(asyncio.gather(*tasks))
-    content_result = list(map(lambda x: x[0], results))
-    id_result = list(map(lambda x: x[1], results))
-    score_result = list(map(lambda x: x[2], results))
+    nested_list = [list(map(lambda x: [f'Query: {query} Document: {x}'], content_list))
+                   for query, content_list in zip(queries, contents_list)]
+
+    rerank_scores = flatten_apply(monot5_run_model, nested_list, model=model, batch_size=batch, tokenizer=tokenizer,
+                                  device=device, token_false_id=token_false_id, token_true_id=token_true_id)
+
+    df = pd.DataFrame({
+        'contents': contents_list,
+        'ids': ids_list,
+        'scores': rerank_scores,
+    })
+    df[['contents', 'ids', 'scores']] = df.apply(sort_by_scores, axis=1, result_type='expand')
+    results = select_top_k(df, ['contents', 'ids', 'scores'], top_k)
 
     del model
     del tokenizer
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
 
-    return content_result, id_result, score_result
-
-
-async def mono_t5_pure(query: str, contents: List[str], scores: List[float], top_k: int,
-                       ids: List[str], model, device, tokenizer, token_false_id, token_true_id)\
-        -> Tuple[List[str], List[str], List[float]]:
-    """
-    Rerank a list of contents based on their relevance to a query using MonoT5.
-
-    :param query: The query to use for reranking
-    :param contents: The list of contents to rerank
-    :param scores: The list of scores retrieved from the initial ranking
-    :param ids: The list of ids retrieved from the initial ranking
-    :param model: The MonoT5 model to use for reranking
-    :param device: The device to run the model on (GPU if available, otherwise CPU)
-    :param tokenizer: The tokenizer to use for the model
-    :param token_false_id: The id of the token used by the model to represent a false prediction
-    :param token_true_id: The id of the token used by the model to represent a true prediction
-    :return: tuple of lists containing the reranked contents, ids, and scores
-    """
-
-    # Format the input for the model by combining each content with the query
-    input_texts = [f'Query: {query} Document: {content}' for content in contents]
-    # Tokenize the input texts and prepare for model input
-    input_encodings = tokenizer(input_texts, padding=True, truncation=True, max_length=512, return_tensors='pt').to(
-        device)
-
-    # Generate model predictions without updating model weights
-    with torch.no_grad():
-        outputs = model.generate(input_ids=input_encodings['input_ids'],
-                                 attention_mask=input_encodings['attention_mask'],
-                                 output_scores=True,
-                                 return_dict_in_generate=True)
-
-    # Extract logits for the 'false' and 'true' tokens from the model's output
-    logits = outputs.scores[-1][:, [token_false_id, token_true_id]]
-    # Calculate the softmax probability of the 'true' token
-    probs = torch.nn.functional.softmax(logits, dim=-1)[:, 1]  # Get the probability of the 'true' token
-
-    # Create a list of tuples pairing each content with its relevance probability
-    content_ids_probs = list(zip(contents, ids, probs.tolist()))
-
-    # Sort the list of pairs based on the relevance score in descending order
-    sorted_content_ids_probs = sorted(content_ids_probs, key=lambda x: x[2], reverse=True)
-
-    # crop with top_k
-    if len(contents) < top_k:
-        top_k = len(contents)
-    sorted_content_ids_probs = sorted_content_ids_probs[:top_k]
-
-    content_result, id_result, score_result = zip(*sorted_content_ids_probs)
-
-    return list(content_result), list(id_result), list(score_result)
+    return results['contents'].tolist(), results['ids'].tolist(), results['scores'].tolist()
+
+
+def monot5_run_model(input_texts, model, batch_size: int, tokenizer, device, token_false_id, token_true_id):
+    batch_input_texts = make_batch(input_texts, batch_size)
+    results = []
+    for batch_texts in batch_input_texts:
+        flattened_batch_texts = list(chain.from_iterable(batch_texts))
+        input_encodings = tokenizer(flattened_batch_texts, padding=True, truncation=True, max_length=512,
+                                    return_tensors='pt').to(
+            device)
+        with torch.no_grad():
+            outputs = model.generate(input_ids=input_encodings['input_ids'],
+                                     attention_mask=input_encodings['attention_mask'],
+                                     output_scores=True,
+                                     return_dict_in_generate=True)
+
+        # Extract logits for the 'false' and 'true' tokens from the model's output
+        logits = outputs.scores[-1][:, [token_false_id, token_true_id]]
+        # Calculate the softmax probability of the 'true' token
+        probs = torch.nn.functional.softmax(logits, dim=-1)[:, 1]
+        results.extend(probs.tolist())
+    return results

autorag/nodes/passagereranker/tart/tart.py

@@ -1,19 +1,21 @@
-import asyncio
+from itertools import chain
 from typing import List, Tuple
 
+import pandas as pd
 import torch
 import torch.nn.functional as F
 
 from autorag.nodes.passagereranker.base import passage_reranker_node
 from autorag.nodes.passagereranker.tart.modeling_enc_t5 import EncT5ForSequenceClassification
 from autorag.nodes.passagereranker.tart.tokenization_enc_t5 import EncT5Tokenizer
+from autorag.utils.util import make_batch, sort_by_scores, flatten_apply, select_top_k
 
 
 @passage_reranker_node
 def tart(queries: List[str], contents_list: List[List[str]],
          scores_list: List[List[float]], ids_list: List[List[str]],
-         top_k: int, instruction: str = "Find passage to answer given question") \
-        -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+         top_k: int, instruction: str = "Find passage to answer given question",
+         batch: int = 64) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
     """
     Rerank a list of contents based on their relevance to a query using Tart.
     TART is a reranker based on TART (https://github.com/facebookresearch/tart).
@@ -29,65 +31,48 @@ def tart(queries: List[str], contents_list: List[List[str]],
         Note: default instruction is "Find passage to answer given question"
             The default instruction from the TART paper is being used.
             If you want to use a different instruction, you can change the instruction through this parameter
+    :param batch: The number of queries to be processed in a batch
     :return: tuple of lists containing the reranked contents, ids, and scores
     """
     model_name = "facebook/tart-full-flan-t5-xl"
     model = EncT5ForSequenceClassification.from_pretrained(model_name)
     tokenizer = EncT5Tokenizer.from_pretrained(model_name)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     model = model.to(device)
-    # Run async tart_rerank_pure function
-    tasks = [tart_pure(query, contents, scores, ids, top_k, model, tokenizer, instruction, device) \
-             for query, contents, scores, ids in zip(queries, contents_list, scores_list, ids_list)]
-    loop = asyncio.get_event_loop()
-    results = loop.run_until_complete(asyncio.gather(*tasks))
-    content_result = list(map(lambda x: x[0], results))
-    id_result = list(map(lambda x: x[1], results))
-    score_result = list(map(lambda x: x[2], results))
+
+    nested_list = [[['{} [SEP] {}'.format(instruction, query)] for _ in contents] for query, contents in
+                   zip(queries, contents_list)]
+
+    rerank_scores = flatten_apply(tart_run_model, nested_list, model=model, batch_size=batch,
+                                  tokenizer=tokenizer, device=device, contents_list=contents_list)
+
+    df = pd.DataFrame({
+        'contents': contents_list,
+        'ids': ids_list,
+        'scores': rerank_scores,
+    })
+    df[['contents', 'ids', 'scores']] = df.apply(sort_by_scores, axis=1, result_type='expand')
+    results = select_top_k(df, ['contents', 'ids', 'scores'], top_k)
 
     del model
     del tokenizer
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
 
-    return content_result, id_result, score_result
-
-
-async def tart_pure(query: str, contents: List[str], scores: List[float],
-                    ids: List[str], top_k: int, model, tokenizer, instruction: str, device) \
-        -> Tuple[List[str], List[str], List[float]]:
-    """
-    Rerank a list of contents based on their relevance to a query using Tart.
-
-    :param query: The query to use for reranking
-    :param contents: The list of contents to rerank
-    :param scores: The list of scores retrieved from the initial ranking
-    :param ids: The list of ids retrieved from the initial ranking
-    :param top_k: The number of passages to be retrieved
-    :param model: The Tart model to use for reranking
-    :param tokenizer: The tokenizer to use for the model
-    :param instruction: The instruction for reranking.
-    :param device: The device to run the model on (GPU if available, otherwise CPU)
-    :return: tuple of lists containing the reranked contents, ids, and scores
-    """
-
-    instruction_queries: List[str] = ['{0} [SEP] {1}'.format(instruction, query) for _ in range(len(contents))]
-    features = tokenizer(instruction_queries, contents, padding=True, truncation=True, return_tensors="pt")
-    features = features.to(device)
-
-    with torch.no_grad():
-        scores = model(**features).logits
-        normalized_scores = [float(score[1]) for score in F.softmax(scores, dim=1)]
-
-    contents_ids_scores = list(zip(contents, ids, normalized_scores))
-
-    sorted_contents_ids_scores = sorted(contents_ids_scores, key=lambda x: x[2], reverse=True)
-
-    # crop with top_k
-    if len(contents) < top_k:
-        top_k = len(contents)
-    sorted_contents_ids_scores = sorted_contents_ids_scores[:top_k]
-
-    content_result, id_result, score_result = zip(*sorted_contents_ids_scores)
-
-    return list(content_result), list(id_result), list(score_result)
+    return results['contents'].tolist(), results['ids'].tolist(), results['scores'].tolist()
+
+
+def tart_run_model(input_texts, contents_list, model, batch_size: int, tokenizer, device):
+    batch_input_texts = make_batch(input_texts, batch_size)
+    batch_contents_list = make_batch(contents_list, batch_size)
+    results = []
+    for batch_texts, batch_contents in zip(batch_input_texts, batch_contents_list):
+        flattened_batch_texts = list(chain.from_iterable(batch_texts))
+        flattened_batch_contents = list(chain.from_iterable(batch_contents))
+        feature = tokenizer(flattened_batch_texts, flattened_batch_contents, padding=True, truncation=True,
+                            return_tensors="pt").to(device)
+        with torch.no_grad():
+            pred_scores = model(**feature).logits
+            normalized_scores = [float(score[1]) for score in F.softmax(pred_scores, dim=1)]
+        results.extend(normalized_scores)
+    return results