Online Retriever

rankify.retrievers.online_retriever

Document

Represents a document consisting of a question, answers, and contexts.

Attributes:

Name	Type	Description
question	Question	The question associated with the document.
answers	Answer	The answers to the question.
contexts	list[Context]	A list of related contexts.
reorder_contexts	list[Context] or None	A reordered list of contexts based on relevance.

Source code in rankify/dataset/dataset.py

class Document:
    """
    Represents a document consisting of a question, answers, and contexts.

    Attributes:
        question (Question): The question associated with the document.
        answers (Answer): The answers to the question.
        contexts (list[Context]): A list of related contexts.
        reorder_contexts (list[Context] or None): A reordered list of contexts based on relevance.
    """
    def __init__(self, question: Question, answers: Answer, contexts: list = None , id: int = None) -> None:
        """
        Initializes a Document instance.

        Args:
            question (Question): The question associated with the document.
            answers (Answer): The answers to the question.
            contexts (list[Context], optional): A list of contexts related to the question.

        Example:
            ```python
            q = Question("What is the capital of France?")
            a = Answer(["Paris"])
            c1 = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            c2 = Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany.")
            d = Document(question=q, answers=a, contexts=[c1, c2])
            print(d)
            ```
        """
        self.question: Question = question
        self.answers: Answer = answers
        self.contexts: List[Context] = contexts
        self.reorder_contexts: List[Context] = None
        self.id = str(id) 

    @classmethod
    def from_dict(cls, data: dict,n_docs:int=100) -> 'Document':
        """
        Creates a Document instance from a dictionary.

        Args:
            data (dict): A dictionary containing the question, answers, and contexts.
            n_docs (int, optional): The number of contexts to include. Defaults to 100.

        Returns:
            Document: A new Document instance.

        Example:
            ```python
            data = {
                "question": "What is the capital of France?",
                "answers": ["Paris"],
                "ctxs": [
                    {"score": 0.9, "has_answer": True, "id": 1, "title": "Paris", "text": "The capital of France is Paris."},
                    {"score": 0.5, "has_answer": False, "id": 2, "title": "Berlin", "text": "Berlin is the capital of Germany."}
                ]
            }
            d = Document.from_dict(data)
            print(d.question)
            ```
        """
        question = Question(data["question"])
        if "answers" in data:
            answers = Answer(data["answers"])
        else:
            answers =Answer('')

        if "query_id" in data:
            id = data["query_id"]
        else:
            id = None
        contexts = [Context(**ctx) for ctx in data["ctxs"][:n_docs]]
        return cls(question, answers, contexts, id=id)

    def to_dict(self) -> Dict[str, Optional[object]]:
        """
        Converts the document into a dictionary representation.

        Returns:
            dict: A dictionary containing the question, answers, and contexts.
        """
        return {
            "question": self.question.question,
            "answers": self.answers.answers,
            "contexts": [ctx.to_dict() for ctx in self.contexts]
        }
    def to_dict_reoreder(self) -> Dict[str,Optional[object]]:
        return {
            "question" : self.question.question,
            "answers" : self.answers.answers,
            "contexts" : [ctx.to_dict() for ctx in self.reorder_contexts]
        }
    def __str__(self) -> str:
        """
        Returns a string representation of the Document instance.

        Returns:
            str: The formatted document information.

        Example:
            ```python
            d = Document(Question("What is the capital of France?"), Answer(["Paris"]))
            print(d)
            ```
        """
        contexts_str = "\n\n".join([str(ctx) for ctx in self.contexts])
        reorder_contexts_str= ''
        if self.reorder_contexts is not None:
            reorder_contexts_str = "\n\n".join([str(ctx) for ctx in self.reorder_contexts])
        return f"{self.question}\n\n{self.answers}\n\nContext: \n\n{contexts_str}\nReorder contexts: \n\n{reorder_contexts_str}"

__init__(question, answers, contexts=None, id=None)

Initializes a Document instance.

Parameters:

Name	Type	Description	Default
question	Question	The question associated with the document.	required
answers	Answer	The answers to the question.	required
contexts	list[Context]	A list of contexts related to the question.	None

Example

q = Question("What is the capital of France?")
a = Answer(["Paris"])
c1 = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
c2 = Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany.")
d = Document(question=q, answers=a, contexts=[c1, c2])
print(d)

Source code in rankify/dataset/dataset.py

def __init__(self, question: Question, answers: Answer, contexts: list = None , id: int = None) -> None:
    """
    Initializes a Document instance.

    Args:
        question (Question): The question associated with the document.
        answers (Answer): The answers to the question.
        contexts (list[Context], optional): A list of contexts related to the question.

    Example:
        ```python
        q = Question("What is the capital of France?")
        a = Answer(["Paris"])
        c1 = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        c2 = Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany.")
        d = Document(question=q, answers=a, contexts=[c1, c2])
        print(d)
        ```
    """
    self.question: Question = question
    self.answers: Answer = answers
    self.contexts: List[Context] = contexts
    self.reorder_contexts: List[Context] = None
    self.id = str(id) 

from_dict(data, n_docs=100) classmethod

Creates a Document instance from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict	A dictionary containing the question, answers, and contexts.	required
n_docs	int	The number of contexts to include. Defaults to 100.	100

Returns:

Name	Type	Description
Document	Document	A new Document instance.

Example

data = {
    "question": "What is the capital of France?",
    "answers": ["Paris"],
    "ctxs": [
        {"score": 0.9, "has_answer": True, "id": 1, "title": "Paris", "text": "The capital of France is Paris."},
        {"score": 0.5, "has_answer": False, "id": 2, "title": "Berlin", "text": "Berlin is the capital of Germany."}
    ]
}
d = Document.from_dict(data)
print(d.question)

Source code in rankify/dataset/dataset.py

@classmethod
def from_dict(cls, data: dict,n_docs:int=100) -> 'Document':
    """
    Creates a Document instance from a dictionary.

    Args:
        data (dict): A dictionary containing the question, answers, and contexts.
        n_docs (int, optional): The number of contexts to include. Defaults to 100.

    Returns:
        Document: A new Document instance.

    Example:
        ```python
        data = {
            "question": "What is the capital of France?",
            "answers": ["Paris"],
            "ctxs": [
                {"score": 0.9, "has_answer": True, "id": 1, "title": "Paris", "text": "The capital of France is Paris."},
                {"score": 0.5, "has_answer": False, "id": 2, "title": "Berlin", "text": "Berlin is the capital of Germany."}
            ]
        }
        d = Document.from_dict(data)
        print(d.question)
        ```
    """
    question = Question(data["question"])
    if "answers" in data:
        answers = Answer(data["answers"])
    else:
        answers =Answer('')

    if "query_id" in data:
        id = data["query_id"]
    else:
        id = None
    contexts = [Context(**ctx) for ctx in data["ctxs"][:n_docs]]
    return cls(question, answers, contexts, id=id)

to_dict()

Converts the document into a dictionary representation.

Returns:

Name	Type	Description
dict	Dict[str, Optional[object]]	A dictionary containing the question, answers, and contexts.

Source code in rankify/dataset/dataset.py

def to_dict(self) -> Dict[str, Optional[object]]:
    """
    Converts the document into a dictionary representation.

    Returns:
        dict: A dictionary containing the question, answers, and contexts.
    """
    return {
        "question": self.question.question,
        "answers": self.answers.answers,
        "contexts": [ctx.to_dict() for ctx in self.contexts]
    }

__str__()

Returns a string representation of the Document instance.

Returns:

Name	Type	Description
str	str	The formatted document information.

Example

d = Document(Question("What is the capital of France?"), Answer(["Paris"]))
print(d)

Source code in rankify/dataset/dataset.py

def __str__(self) -> str:
    """
    Returns a string representation of the Document instance.

    Returns:
        str: The formatted document information.

    Example:
        ```python
        d = Document(Question("What is the capital of France?"), Answer(["Paris"]))
        print(d)
        ```
    """
    contexts_str = "\n\n".join([str(ctx) for ctx in self.contexts])
    reorder_contexts_str= ''
    if self.reorder_contexts is not None:
        reorder_contexts_str = "\n\n".join([str(ctx) for ctx in self.reorder_contexts])
    return f"{self.question}\n\n{self.answers}\n\nContext: \n\n{contexts_str}\nReorder contexts: \n\n{reorder_contexts_str}"

Context

Represents a context with metadata such as score and title.

Attributes:

Name	Type	Description
score	float	The relevance score of the context.
has_answer	bool	Whether the context contains an answer.
id	int	The identifier of the context.
title	str	The title of the context.
text	str	The text of the context.

Source code in rankify/dataset/dataset.py

class Context:
    """
    Represents a context with metadata such as score and title.

    Attributes:
        score (float, optional): The relevance score of the context.
        has_answer (bool, optional): Whether the context contains an answer.
        id (int, optional): The identifier of the context.
        title (str, optional): The title of the context.
        text (str, optional): The text of the context.
    """
    def __init__(self, score: float=None, has_answer: bool=None, id: str=None, title: str=None, text: str=None)-> None:
        """
        Initializes a Context instance.

        Args:
            score (float, optional): The relevance score.
            has_answer (bool, optional): Whether the context contains an answer.
            id (int, optional): The identifier of the context.
            title (str, optional): The title of the context.
            text (str, optional): The text of the context.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(c)
            ```
        """
        self.score: Optional[float] = score
        self.has_answer: Optional[bool] = has_answer
        self.id: Optional[str] = id
        self.title: Optional[str] = title
        self.text: Optional[str] = text

    def to_dict(self, save_text: bool=False) -> Dict[str, Optional[object]]:

        """
        Converts the Context instance to a dictionary.

        Args:
            save_text (bool): Whether to include text in the output dictionary.

        Returns:
            dict: The context data.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(c.to_dict())
            ```
        """
        context_dict = {
            "score": float(self.score) if self.score is not None else None,
            "has_answer": self.has_answer,
            "id": self.id,
            }

        # Include 'text' only if save_text is True
        if save_text:
            context_dict["text"] = self.text
            context_dict["title"] =  self.title

        return context_dict
    def __str__(self) -> str:
        """
        Returns a string representation of the Context instance.

        Returns:
            str: The formatted context.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(str(c))
            ```
        """
        return f"ID: {self.id}\nHas Answer: {self.has_answer}\nTitle: {self.title}\nText: {self.text}\nScore: {self.score}"

__init__(score=None, has_answer=None, id=None, title=None, text=None)

Initializes a Context instance.

Parameters:

Name	Type	Description	Default
score	float	The relevance score.	None
has_answer	bool	Whether the context contains an answer.	None
id	int	The identifier of the context.	None
title	str	The title of the context.	None
text	str	The text of the context.	None

Example

c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(c)

Source code in rankify/dataset/dataset.py

def __init__(self, score: float=None, has_answer: bool=None, id: str=None, title: str=None, text: str=None)-> None:
    """
    Initializes a Context instance.

    Args:
        score (float, optional): The relevance score.
        has_answer (bool, optional): Whether the context contains an answer.
        id (int, optional): The identifier of the context.
        title (str, optional): The title of the context.
        text (str, optional): The text of the context.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(c)
        ```
    """
    self.score: Optional[float] = score
    self.has_answer: Optional[bool] = has_answer
    self.id: Optional[str] = id
    self.title: Optional[str] = title
    self.text: Optional[str] = text

to_dict(save_text=False)

Converts the Context instance to a dictionary.

Parameters:

Name	Type	Description	Default
save_text	bool	Whether to include text in the output dictionary.	False

Returns:

Name	Type	Description
dict	Dict[str, Optional[object]]	The context data.

Example

c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(c.to_dict())

Source code in rankify/dataset/dataset.py

def to_dict(self, save_text: bool=False) -> Dict[str, Optional[object]]:

    """
    Converts the Context instance to a dictionary.

    Args:
        save_text (bool): Whether to include text in the output dictionary.

    Returns:
        dict: The context data.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(c.to_dict())
        ```
    """
    context_dict = {
        "score": float(self.score) if self.score is not None else None,
        "has_answer": self.has_answer,
        "id": self.id,
        }

    # Include 'text' only if save_text is True
    if save_text:
        context_dict["text"] = self.text
        context_dict["title"] =  self.title

    return context_dict

__str__()

Returns a string representation of the Context instance.

Returns:

Name	Type	Description
str	str	The formatted context.

Example

c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(str(c))

Source code in rankify/dataset/dataset.py

def __str__(self) -> str:
    """
    Returns a string representation of the Context instance.

    Returns:
        str: The formatted context.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(str(c))
        ```
    """
    return f"ID: {self.id}\nHas Answer: {self.has_answer}\nTitle: {self.title}\nText: {self.text}\nScore: {self.score}"

WebSearchTool

Bases: Tool

Source code in rankify/tools/Tools.py

class WebSearchTool(Tool):

    name = "web_search"
    description = "Performs web search using user query and returns the top n results."
    inputs = {
        'query': {
            'type': 'string',
            'description': 'The user search query'
        }
    }
    output_type = "string"

    def __init__(self, model: Optional[str]=None, search_provider: str = 'serper',
                 search_provider_api_key: Optional[str] = None):
        super().__init__()
        self.search_client = None
        self.model = model
        self.search_provider = search_provider
        self.search_provider_api_key = search_provider_api_key
#        self.validate_arguments()

    @staticmethod
    def create_search_api(search_provider: str = 'SERPER', api_key: Optional[str] = None) -> SearchAPIClient:
        """
        Instantiate a search API client to fetch SERP results from a SERP provider e.g., SERPER.dev.
        Args:
            search_provider (str): The name of the search provider.
            Api_key (str): The API key for the search provider. Can be stored and loaded from env vars.
        Returns:
            SearchAPIClient: An instance of the search API client. E.g., SerpAPIClient.
        """
        if search_provider.lower() == 'serper':
            #print(api_key)
            return SerperApiClient(api_key=api_key)
        else:
            raise ValueError(f'Invalid search provider{search_provider}')

    def forward(self, query: str,num_result:int=10):
        sources = self.search_client.search_web(query=query,num_results=num_result)
        #print(sources, len(sources.data['organic']), num_result)


        #asdadsasdad
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                nest_asyncio.apply()
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)

        return loop.run_until_complete(build_search_context(sources))



    def setup(self):
        """
        Performs time-consuming setup operations here e.g., model loading.
        """
        self.search_client = self.create_search_api(api_key=self.search_provider_api_key)
        self.is_initialized = True

create_search_api(search_provider='SERPER', api_key=None) staticmethod

Instantiate a search API client to fetch SERP results from a SERP provider e.g., SERPER.dev. Args: search_provider (str): The name of the search provider. Api_key (str): The API key for the search provider. Can be stored and loaded from env vars. Returns: SearchAPIClient: An instance of the search API client. E.g., SerpAPIClient.

Source code in rankify/tools/Tools.py

@staticmethod
def create_search_api(search_provider: str = 'SERPER', api_key: Optional[str] = None) -> SearchAPIClient:
    """
    Instantiate a search API client to fetch SERP results from a SERP provider e.g., SERPER.dev.
    Args:
        search_provider (str): The name of the search provider.
        Api_key (str): The API key for the search provider. Can be stored and loaded from env vars.
    Returns:
        SearchAPIClient: An instance of the search API client. E.g., SerpAPIClient.
    """
    if search_provider.lower() == 'serper':
        #print(api_key)
        return SerperApiClient(api_key=api_key)
    else:
        raise ValueError(f'Invalid search provider{search_provider}')

setup()

Performs time-consuming setup operations here e.g., model loading.

Source code in rankify/tools/Tools.py

def setup(self):
    """
    Performs time-consuming setup operations here e.g., model loading.
    """
    self.search_client = self.create_search_api(api_key=self.search_provider_api_key)
    self.is_initialized = True

Chunker

A modular text chunking class that splits text into smaller, overlapping segments.

This class provides a flexible way to break down large texts into smaller chunks while maintaining context through configurable overlap. It uses RecursiveCharacterTextSplitter from langchain under the hood.

Attributes:

Name	Type	Description
chunk_size	int	The target size for each text chunk.
chunk_overlap	int	The number of characters to overlap between chunks.
separators	List[str]	List of separators to use for splitting, in order of preference.
length_function	callable	Function to measure text length (default: len).

Source code in rankify/tools/websearch/context/chunker.py

class Chunker:
    """A modular text chunking class that splits text into smaller, overlapping segments.

    This class provides a flexible way to break down large texts into smaller chunks
    while maintaining context through configurable overlap. It uses RecursiveCharacterTextSplitter
    from langchain under the hood.

    Attributes:
        chunk_size (int): The target size for each text chunk.
        chunk_overlap (int): The number of characters to overlap between chunks.
        separators (List[str]): List of separators to use for splitting, in order of preference.
        length_function (callable): Function to measure text length (default: len).
    """

    def __init__(
        self,
        chunk_size: int = 1200,
        chunk_overlap: int = 100,
        min_chunk_size: int = 800, 
        separators: Optional[List[str]] = None,
        length_function: callable = len
    ):
        """Initialize the Chunker with specified parameters.

        Args:
            chunk_size (int, optional): Target size for each chunk. Defaults to 250.
            chunk_overlap (int, optional): Number of characters to overlap. Defaults to 50.
            separators (List[str], optional): Custom separators for splitting.
                Defaults to ["\n\n", "\n", " "].
            length_function (callable, optional): Function to measure text length.
                Defaults to len.
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.separators = separators or ["\n\n", "\n",]
        self.length_function = length_function
        self.min_chunk_size = min_chunk_size

        self.splitter = RecursiveCharacterTextSplitter(
            separators=self.separators,
            chunk_size=self.chunk_size,
            chunk_overlap=self.chunk_overlap,
            length_function=self.length_function
        )

    def split_text(self, text: str) -> List[str]:
        """Split text and combine small chunks."""
        chunks = self.splitter.split_text(text)

        # Combine chunks that are too small
        combined_chunks = []
        current_chunk = ""

        for chunk in chunks:
            chunk = self.remove_links(chunk)
            if len(current_chunk) + len(chunk) < self.min_chunk_size:
                current_chunk += (" " if current_chunk else "") + chunk
            else:
                if current_chunk:
                    combined_chunks.append(current_chunk)
                current_chunk = chunk

        if current_chunk:
            combined_chunks.append(current_chunk)

        return combined_chunks
    def remove_links(self, text: str) -> str:
        """Remove various types of links from text."""
        patterns = [
            r'https?://[^\s<>"{}|\\^`\[\]]+',  # HTTP/HTTPS URLs
            r'www\.[^\s<>"{}|\\^`\[\]]+',      # www. links
            r'ftp://[^\s<>"{}|\\^`\[\]]+',     # FTP links
            r'\[[^\]]*\]\([^)]*\)',            # Markdown links [text](url)
            r'<a[^>]*>.*?</a>',                # HTML links
        ]

        for pattern in patterns:
            text = re.sub(pattern, '', text, flags=re.IGNORECASE)

        # Clean up extra whitespace
        text = re.sub(r'\s+', ' ', text).strip()
        return text
    def split_texts(self, texts: List[str]) -> List[List[str]]:
        """Split multiple texts into chunks.

        Args:
            texts (List[str]): A list of input texts to be split into chunks.

        Returns:
            List[List[str]]: A list of lists, where each inner list contains
                the chunks for one input text.
        """
        return [self.split_text(text) for text in texts]

__init__(chunk_size=1200, chunk_overlap=100, min_chunk_size=800, separators=None, length_function=len)

Initialize the Chunker with specified parameters.

    Args:
        chunk_size (int, optional): Target size for each chunk. Defaults to 250.
        chunk_overlap (int, optional): Number of characters to overlap. Defaults to 50.
        separators (List[str], optional): Custom separators for splitting.
            Defaults to ["

", " ", " "]. length_function (callable, optional): Function to measure text length. Defaults to len.

Source code in rankify/tools/websearch/context/chunker.py

def __init__(
    self,
    chunk_size: int = 1200,
    chunk_overlap: int = 100,
    min_chunk_size: int = 800, 
    separators: Optional[List[str]] = None,
    length_function: callable = len
):
    """Initialize the Chunker with specified parameters.

    Args:
        chunk_size (int, optional): Target size for each chunk. Defaults to 250.
        chunk_overlap (int, optional): Number of characters to overlap. Defaults to 50.
        separators (List[str], optional): Custom separators for splitting.
            Defaults to ["\n\n", "\n", " "].
        length_function (callable, optional): Function to measure text length.
            Defaults to len.
    """
    self.chunk_size = chunk_size
    self.chunk_overlap = chunk_overlap
    self.separators = separators or ["\n\n", "\n",]
    self.length_function = length_function
    self.min_chunk_size = min_chunk_size

    self.splitter = RecursiveCharacterTextSplitter(
        separators=self.separators,
        chunk_size=self.chunk_size,
        chunk_overlap=self.chunk_overlap,
        length_function=self.length_function
    )

split_text(text)

Split text and combine small chunks.

Source code in rankify/tools/websearch/context/chunker.py

def split_text(self, text: str) -> List[str]:
    """Split text and combine small chunks."""
    chunks = self.splitter.split_text(text)

    # Combine chunks that are too small
    combined_chunks = []
    current_chunk = ""

    for chunk in chunks:
        chunk = self.remove_links(chunk)
        if len(current_chunk) + len(chunk) < self.min_chunk_size:
            current_chunk += (" " if current_chunk else "") + chunk
        else:
            if current_chunk:
                combined_chunks.append(current_chunk)
            current_chunk = chunk

    if current_chunk:
        combined_chunks.append(current_chunk)

    return combined_chunks

remove_links(text)

Remove various types of links from text.

Source code in rankify/tools/websearch/context/chunker.py

def remove_links(self, text: str) -> str:
    """Remove various types of links from text."""
    patterns = [
        r'https?://[^\s<>"{}|\\^`\[\]]+',  # HTTP/HTTPS URLs
        r'www\.[^\s<>"{}|\\^`\[\]]+',      # www. links
        r'ftp://[^\s<>"{}|\\^`\[\]]+',     # FTP links
        r'\[[^\]]*\]\([^)]*\)',            # Markdown links [text](url)
        r'<a[^>]*>.*?</a>',                # HTML links
    ]

    for pattern in patterns:
        text = re.sub(pattern, '', text, flags=re.IGNORECASE)

    # Clean up extra whitespace
    text = re.sub(r'\s+', ' ', text).strip()
    return text

split_texts(texts)

Split multiple texts into chunks.

Parameters:

Name	Type	Description	Default
texts	List[str]	A list of input texts to be split into chunks.	required

Returns:

Type	Description
List[List[str]]	List[List[str]]: A list of lists, where each inner list contains the chunks for one input text.

Source code in rankify/tools/websearch/context/chunker.py

def split_texts(self, texts: List[str]) -> List[List[str]]:
    """Split multiple texts into chunks.

    Args:
        texts (List[str]): A list of input texts to be split into chunks.

    Returns:
        List[List[str]]: A list of lists, where each inner list contains
            the chunks for one input text.
    """
    return [self.split_text(text) for text in texts]

OnlineRetriever

Source code in rankify/retrievers/online_retriever.py

class OnlineRetriever:
    def __init__(
        self,
        n_docs: int = 5,
        api_key: str = None,
        chunk_size: int = 500,
        chunk_overlap: int = 50,
        **kwags
    ) -> None:
        self.n_docs = n_docs
        self.searcher = WebSearchTool(search_provider_api_key=api_key)
        self.tokenizer = SimpleTokenizer()
        self.chunker = Chunker(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
        )

    def _search_web(self, query: str) -> List[Any]:
        if not self.searcher.is_initialized:
            self.searcher.setup()
        return self.searcher.forward(query, num_result=self.n_docs)


    def retrieve(self, documents: List[Document]) -> List[Document]:
        for idx, doc in enumerate(tqdm(documents, desc="Fetching docs...")):
            question = doc.question.question
            logger.info(f"Fetching contexts for Q{idx}: {question}")

            sources = self._search_web(question)
            processed = []
            for s_idx, source in enumerate(sources):
                text = source.get('fit_markdown', '')
                # with open("log.txt", "a",  encoding="utf-8") as f:
                #     f.writelines(html)
                if text:
                    passages = self.chunker.split_text(text)
                    #print(len(passages), "passages found")
                    for p_idx, p in enumerate(passages):
                        #print(p)

                        #print(f"Processing snippet {s_idx} passage {p_idx}: {p['contents'][:50]}...")
                        processed.append({'id': f'{s_idx}_{p_idx}', 'contents': p})
                else:
                    snippet = source.get('snippet','')
                    if len(snippet.split()) > 20:
                        processed.append({'id': f's{ s_idx }_snip', 'contents': snippet})

            # Indexing
            with tempfile.TemporaryDirectory() as tmpdir:
                json.dump(processed, open(os.path.join(tmpdir,'docs.json'),'w'))
                subprocess.run([
                    'python','-m','pyserini.index.lucene',
                    '-collection','JsonCollection','-generator','DefaultLuceneDocumentGenerator',
                    '-input',tmpdir,'-index',tmpdir,
                    '-storePositions','-storeDocvectors','-storeRaw'
                ], check=True)

                searcher = LuceneSearcher(tmpdir)
                hits = searcher.search(question, k=self.n_docs)
                contexts: List[Context] = []
                for h_idx, hit in enumerate(hits):
                    raw = searcher.doc(hit.docid).raw()
                    data = json.loads(raw)
                    text = data.get('contents','')
                    title = "" #text.split("\n")[0]
                    try:
                        cid = int(hit.docid)
                    except ValueError:
                        cid = h_idx
                    contexts.append(Context(
                        id=cid,
                        title=title,
                        text=text,
                        score=hit.score,
                        has_answer=has_answers(text, doc.answers.answers, self.tokenizer)
                    ))
                doc.contexts = contexts
        return documents