Skip to content

code_parser

langroid/parsing/code_parser.py

CodeParser(config)

Source code in langroid/parsing/code_parser.py 
75
76
77
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 
79
80
81
82
83
84
85
86
87
88
89
90
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. Each chunk gets is_chunk=True, a distinct id, and the ids of its neighbors within the same source document -- so the result can be given straight to a vector store, as with the splitters on Parser. Calling Parser.add_window_ids on the result afterwards is a no-op and is no longer needed.

Source code in langroid/parsing/code_parser.py 
 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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
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. Each chunk
        gets `is_chunk=True`, a distinct id, and the ids of its neighbors within
        the same source document -- so the result can be given straight to a
        vector store, as with the splitters on `Parser`. Calling
        `Parser.add_window_ids` on the result afterwards is a no-op and is no
        longer needed.
    """
    chunked_docs = [
        [
            Document(
                content=chunk,
                # window_ids is cleared, not inherited: a chunk's window
                # names its own siblings, so carrying over the window of
                # an already-chunked input would both be wrong and make
                # these look already-windowed to add_window_ids
                metadata=d.metadata.model_copy(
                    update=dict(is_chunk=True, window_ids=[])
                ),
            )
            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
    ]
    if len(chunked_docs) == 0:
        return []
    # per source doc, as in Parser.split_simple: without distinct ids the
    # chunks of one doc overwrite each other on a vector-store upsert,
    # since those are keyed on the id
    for group in chunked_docs:
        add_window_ids(group, self.config.n_neighbor_ids)
    # 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 
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
47
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