Skip to content

code_parser

langroid/parsing/code_parser.py

CodeParser(config)

Source code in langroid/parsing/code_parser.py 
73
74
75
def __init__(self, config: CodeParsingConfig):
    self.config = config
    self.tokenizer = tiktoken.encoding_for_model(config.token_encoding_model)

num_tokens(text)

How many tokens are in the text, according to the tokenizer. This needs to be accurate, otherwise we may exceed the maximum number of tokens allowed by the model. Args: text: string to tokenize Returns: number of tokens in the text

Source code in langroid/parsing/code_parser.py 
77
78
79
80
81
82
83
84
85
86
87
88
def num_tokens(self, text: str) -> int:
    """
    How many tokens are in the text, according to the tokenizer.
    This needs to be accurate, otherwise we may exceed the maximum
    number of tokens allowed by the model.
    Args:
        text: string to tokenize
    Returns:
        number of tokens in the text
    """
    tokens = self.tokenizer.encode(text)
    return len(tokens)

split(docs)

Split the documents into chunks, according to the config.splitter. Only the documents with a language in the config.extensions are split.

Note

We assume the metadata in each document has at least a language field, which is used to determine how to chunk the code.

Args: docs: list of documents to split Returns: list of documents, where each document is a chunk; the metadata of the original document is duplicated for each chunk, so that when we retrieve a chunk, we immediately know info about the original document.

Source code in langroid/parsing/code_parser.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
def split(self, docs: List[Document]) -> List[Document]:
    """
    Split the documents into chunks, according to the config.splitter.
    Only the documents with a language in the config.extensions are split.
    !!! note
        We assume the metadata in each document has at least a `language` field,
        which is used to determine how to chunk the code.
    Args:
        docs: list of documents to split
    Returns:
        list of documents, where each document is a chunk; the metadata of the
        original document is duplicated for each chunk, so that when we retrieve a
        chunk, we immediately know info about the original document.
    """
    chunked_docs = [
        [
            Document(content=chunk, metadata=d.metadata)
            for chunk in chunk_code(
                d.content,
                d.metadata.language,  # type: ignore
                self.config.chunk_size,
                self.num_tokens,
            )
            if chunk.strip() != ""
        ]
        for d in docs
        if d.metadata.language in self.config.extensions  # type: ignore
    ]
    # collapse the list of lists into a single list
    return reduce(lambda x, y: x + y, chunked_docs)

chunk_code(code, language, max_tokens, len_fn)

Chunk code into smaller pieces, so that we don't exceed the maximum number of tokens allowed by the embedding model. Args: code: string of code language: str as a file extension, e.g. "py", "yml" max_tokens: max tokens per chunk len_fn: function to get the length of a string in token units Returns:

Source code in langroid/parsing/code_parser.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def chunk_code(
    code: str, language: str, max_tokens: int, len_fn: Callable[[str], int]
) -> List[str]:
    """
    Chunk code into smaller pieces, so that we don't exceed the maximum
    number of tokens allowed by the embedding model.
    Args:
        code: string of code
        language: str as a file extension, e.g. "py", "yml"
        max_tokens: max tokens per chunk
        len_fn: function to get the length of a string in token units
    Returns:

    """
    lexer = get_lexer_by_name(language)
    tokens = list(lex(code, lexer))

    chunks = []
    current_chunk = ""
    for token_type, token_value in tokens:
        if token_type in Token.Text.Whitespace:
            current_chunk += token_value
        else:
            token_tokens = len_fn(token_value)
            if len_fn(current_chunk) + token_tokens <= max_tokens:
                current_chunk += token_value
            else:
                chunks.append(current_chunk)
                current_chunk = token_value

    if current_chunk:
        chunks.append(current_chunk)

    return chunks