code-graph-rag 0.0.79__py3-none-any.whl

codebase_rag/unixcoder.py

@@ -0,0 +1,278 @@
+# (H) Adapted from https://github.com/microsoft/unixcoder
+# (H) Copyright (c) Microsoft Corporation.
+# (H) Licensed under the MIT license.
+
+import torch
+from torch import nn
+from transformers import RobertaConfig, RobertaModel, RobertaTokenizer
+
+from . import constants as cs
+
+
+class UniXcoder(nn.Module):
+    def __init__(self, model_name: str) -> None:
+        super().__init__()
+        self.tokenizer: RobertaTokenizer = RobertaTokenizer.from_pretrained(model_name)
+        self.config: RobertaConfig = RobertaConfig.from_pretrained(model_name)
+        self.config.is_decoder = True
+        self.model: RobertaModel = RobertaModel.from_pretrained(
+            model_name, config=self.config
+        )
+
+        self.register_buffer(
+            cs.UNIXCODER_BUFFER_BIAS,
+            torch.tril(
+                torch.ones(
+                    (cs.UNIXCODER_MAX_CONTEXT, cs.UNIXCODER_MAX_CONTEXT),
+                    dtype=torch.uint8,
+                )
+            ).view(1, cs.UNIXCODER_MAX_CONTEXT, cs.UNIXCODER_MAX_CONTEXT),
+        )
+        self.lm_head: nn.Linear = nn.Linear(
+            self.config.hidden_size, self.config.vocab_size, bias=False
+        )
+        self.lm_head.weight = self.model.embeddings.word_embeddings.weight
+        self.lsm: nn.LogSoftmax = nn.LogSoftmax(dim=-1)
+
+        self.tokenizer.add_tokens([cs.UNIXCODER_MASK_TOKEN], special_tokens=True)
+
+    def tokenize(
+        self,
+        inputs: list[str],
+        mode: cs.UniXcoderMode = cs.UniXcoderMode.ENCODER_ONLY,
+        max_length: int = 512,
+        padding: bool = False,
+    ) -> list[list[int]]:
+        assert max_length < cs.UNIXCODER_MAX_CONTEXT
+
+        tokenizer = self.tokenizer
+
+        tokens_ids = []
+        for x in inputs:
+            tokens = tokenizer.tokenize(x)
+            match mode:
+                case cs.UniXcoderMode.ENCODER_ONLY:
+                    tokens = tokens[: max_length - 4]
+                    tokens = (
+                        [tokenizer.cls_token, mode, tokenizer.sep_token]
+                        + tokens
+                        + [tokenizer.sep_token]
+                    )
+                case cs.UniXcoderMode.DECODER_ONLY:
+                    tokens = tokens[-(max_length - 3) :]
+                    tokens = [tokenizer.cls_token, mode, tokenizer.sep_token] + tokens
+                case cs.UniXcoderMode.ENCODER_DECODER:
+                    tokens = tokens[: max_length - 5]
+                    tokens = (
+                        [tokenizer.cls_token, mode, tokenizer.sep_token]
+                        + tokens
+                        + [tokenizer.sep_token]
+                    )
+
+            converted = tokenizer.convert_tokens_to_ids(tokens)
+            tokens_id: list[int] = (
+                converted if isinstance(converted, list) else [converted]
+            )
+            if padding:
+                pad_id = self.config.pad_token_id
+                assert pad_id is not None
+                tokens_id += [pad_id] * (max_length - len(tokens_id))
+            tokens_ids.append(tokens_id)
+        return tokens_ids
+
+    def decode(self, source_ids: torch.Tensor) -> list[list[str]]:
+        predictions = []
+        for x in source_ids:
+            prediction = []
+            for y in x:
+                t = y.cpu().numpy()
+                t = list(t)
+                if 0 in t:
+                    t = t[: t.index(0)]
+                text = self.tokenizer.decode(t, clean_up_tokenization_spaces=False)
+                prediction.append(text)
+            predictions.append(prediction)
+        return predictions
+
+    def forward(self, source_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        pad_id = self.config.pad_token_id
+        assert pad_id is not None
+        mask = source_ids.ne(pad_id)
+        token_embeddings = self.model(
+            source_ids, attention_mask=mask.unsqueeze(1) * mask.unsqueeze(2)
+        )[0]
+        sentence_embeddings = (token_embeddings * mask.unsqueeze(-1)).sum(1) / mask.sum(
+            -1
+        ).unsqueeze(-1)
+        return token_embeddings, sentence_embeddings
+
+    def generate(
+        self,
+        source_ids: torch.Tensor,
+        decoder_only: bool = True,
+        eos_id: int | None = None,
+        beam_size: int = 5,
+        max_length: int = 64,
+    ) -> torch.Tensor:
+        # (H) self.bias is registered as buffer (Tensor) but typed as Module by ty
+        bias: torch.Tensor = getattr(self, cs.UNIXCODER_BUFFER_BIAS)
+        pad_id = self.config.pad_token_id
+        assert pad_id is not None
+
+        if decoder_only:
+            mask = bias[:, : source_ids.size(-1), : source_ids.size(-1)]
+        else:
+            mask = source_ids.ne(pad_id)
+            mask = mask.unsqueeze(1) * mask.unsqueeze(2)
+
+        if eos_id is None:
+            eos_id = self.config.eos_token_id
+        assert eos_id is not None
+
+        device = source_ids.device
+
+        preds = []
+        zero = torch.LongTensor(1).fill_(0).to(device)
+        source_len = list(source_ids.ne(1).sum(-1).cpu().numpy())
+        length = source_ids.size(-1)
+        encoder_output = self.model(source_ids, attention_mask=mask)
+        for i in range(source_ids.shape[0]):
+            context = [
+                [x[i : i + 1, :, : source_len[i]].repeat(beam_size, 1, 1, 1) for x in y]
+                for y in encoder_output.past_key_values
+            ]
+            beam = Beam(beam_size, eos_id, device)
+            input_ids = beam.getCurrentState().clone()
+            context_ids = source_ids[i : i + 1, : source_len[i]].repeat(beam_size, 1)
+            out = encoder_output.last_hidden_state[i : i + 1, : source_len[i]].repeat(
+                beam_size, 1, 1
+            )
+            for _ in range(max_length):
+                if beam.done():
+                    break
+                if _ == 0:
+                    hidden_states = out[:, -1, :]
+                    out = self.lsm(self.lm_head(hidden_states)).data
+                    beam.advance(out)
+                    input_ids.data.copy_(
+                        input_ids.data.index_select(0, beam.getCurrentOrigin())
+                    )
+                    input_ids = beam.getCurrentState().clone()
+                else:
+                    length = context_ids.size(-1) + input_ids.size(-1)
+                    out = self.model(
+                        input_ids,
+                        attention_mask=bias[:, context_ids.size(-1) : length, :length],
+                        past_key_values=context,
+                    ).last_hidden_state
+                    hidden_states = out[:, -1, :]
+                    out = self.lsm(self.lm_head(hidden_states)).data
+                    beam.advance(out)
+                    input_ids.data.copy_(
+                        input_ids.data.index_select(0, beam.getCurrentOrigin())
+                    )
+                    input_ids = torch.cat(
+                        (input_ids, beam.getCurrentState().clone()), -1
+                    )
+            hyp = beam.getHyp(beam.getFinal())
+            pred = beam.buildTargetTokens(hyp)[:beam_size]
+            pred = [
+                torch.cat(
+                    [x.view(-1) for x in p] + [zero] * (max_length - len(p))
+                ).view(1, -1)
+                for p in pred
+            ]
+            preds.append(torch.cat(pred, 0).unsqueeze(0))
+
+        preds = torch.cat(preds, 0)
+
+        return preds
+
+
+class Beam:
+    def __init__(self, size: int, eos: int, device: torch.device) -> None:
+        self.size = size
+        self.device = device
+        self.scores: torch.Tensor = torch.FloatTensor(size).zero_().to(device)
+        self.prevKs: list[torch.Tensor] = []
+        self.nextYs: list[torch.Tensor] = [torch.LongTensor(size).fill_(0).to(device)]
+        self._eos = eos
+        self.eosTop = False
+        self.finished: list[tuple[torch.Tensor, int, int]] = []
+
+    def getCurrentState(self) -> torch.Tensor:
+        batch = self.nextYs[-1].view(-1, 1)
+        return batch
+
+    def getCurrentOrigin(self) -> torch.Tensor:
+        return self.prevKs[-1]
+
+    def advance(self, wordLk: torch.Tensor) -> None:
+        numWords = wordLk.size(1)
+
+        if len(self.prevKs) > 0:
+            beamLk = wordLk + self.scores.unsqueeze(1).expand_as(wordLk)
+
+            for i in range(self.nextYs[-1].size(0)):
+                if self.nextYs[-1][i] == self._eos:
+                    beamLk[i] = -1e20
+        else:
+            beamLk = wordLk[0]
+        flatBeamLk = beamLk.view(-1)
+        bestScores, bestScoresId = flatBeamLk.topk(self.size, 0, True, True)
+
+        self.scores = bestScores
+
+        prevK = torch.div(bestScoresId, numWords, rounding_mode="floor")
+        self.prevKs.append(prevK)
+        self.nextYs.append(bestScoresId - prevK * numWords)
+
+        for i in range(self.nextYs[-1].size(0)):
+            if self.nextYs[-1][i] == self._eos:
+                s = self.scores[i]
+                self.finished.append((s, len(self.nextYs) - 1, i))
+
+        if self.nextYs[-1][0] == self._eos:
+            self.eosTop = True
+
+    def done(self) -> bool:
+        return self.eosTop and len(self.finished) >= self.size
+
+    def getFinal(self) -> list[tuple[torch.Tensor, int, int]]:
+        if len(self.finished) == 0:
+            self.finished.append((self.scores[0], len(self.nextYs) - 1, 0))
+        self.finished.sort(key=lambda a: -a[0])
+        if len(self.finished) != self.size:
+            unfinished = [
+                (self.scores[i], len(self.nextYs) - 1, i)
+                for i in range(self.nextYs[-1].size(0))
+                if self.nextYs[-1][i] != self._eos
+            ]
+            unfinished.sort(key=lambda a: -a[0])
+            self.finished += unfinished[: self.size - len(self.finished)]
+        return self.finished[: self.size]
+
+    def getHyp(
+        self, beam_res: list[tuple[torch.Tensor, int, int]]
+    ) -> list[list[torch.Tensor]]:
+        hyps: list[list[torch.Tensor]] = []
+        for _, timestep, k in beam_res:
+            hyp: list[torch.Tensor] = []
+            for j in range(len(self.prevKs[:timestep]) - 1, -1, -1):
+                hyp.append(self.nextYs[j + 1][k])
+                k = self.prevKs[j][k]
+            hyps.append(hyp[::-1])
+        return hyps
+
+    def buildTargetTokens(
+        self, preds: list[list[torch.Tensor]]
+    ) -> list[list[torch.Tensor]]:
+        sentence: list[list[torch.Tensor]] = []
+        for pred in preds:
+            tokens: list[torch.Tensor] = []
+            for tok in pred:
+                if tok == self._eos:
+                    break
+                tokens.append(tok)
+            sentence.append(tokens)
+        return sentence

codebase_rag/vector_store.py

@@ -0,0 +1,80 @@
+from loguru import logger
+
+from . import logs as ls
+from .config import settings
+from .constants import PAYLOAD_NODE_ID, PAYLOAD_QUALIFIED_NAME
+from .utils.dependencies import has_qdrant_client
+
+if has_qdrant_client():
+    from qdrant_client import QdrantClient
+    from qdrant_client.models import Distance, PointStruct, VectorParams
+
+    _CLIENT: QdrantClient | None = None
+
+    def get_qdrant_client() -> QdrantClient:
+        global _CLIENT
+        if _CLIENT is None:
+            _CLIENT = QdrantClient(path=settings.QDRANT_DB_PATH)
+            if not _CLIENT.collection_exists(settings.QDRANT_COLLECTION_NAME):
+                _CLIENT.create_collection(
+                    collection_name=settings.QDRANT_COLLECTION_NAME,
+                    vectors_config=VectorParams(
+                        size=settings.QDRANT_VECTOR_DIM, distance=Distance.COSINE
+                    ),
+                )
+        return _CLIENT
+
+    def store_embedding(
+        node_id: int, embedding: list[float], qualified_name: str
+    ) -> None:
+        try:
+            client = get_qdrant_client()
+            client.upsert(
+                collection_name=settings.QDRANT_COLLECTION_NAME,
+                points=[
+                    PointStruct(
+                        id=node_id,
+                        vector=embedding,
+                        payload={
+                            PAYLOAD_NODE_ID: node_id,
+                            PAYLOAD_QUALIFIED_NAME: qualified_name,
+                        },
+                    )
+                ],
+            )
+        except Exception as e:
+            logger.warning(
+                ls.EMBEDDING_STORE_FAILED.format(name=qualified_name, error=e)
+            )
+
+    def search_embeddings(
+        query_embedding: list[float], top_k: int | None = None
+    ) -> list[tuple[int, float]]:
+        effective_top_k = top_k if top_k is not None else settings.QDRANT_TOP_K
+        try:
+            client = get_qdrant_client()
+            result = client.query_points(
+                collection_name=settings.QDRANT_COLLECTION_NAME,
+                query=query_embedding,
+                limit=effective_top_k,
+            )
+            return [
+                (hit.payload[PAYLOAD_NODE_ID], hit.score)
+                for hit in result.points
+                if hit.payload is not None
+            ]
+        except Exception as e:
+            logger.warning(ls.EMBEDDING_SEARCH_FAILED.format(error=e))
+            return []
+
+else:
+
+    def store_embedding(
+        node_id: int, embedding: list[float], qualified_name: str
+    ) -> None:
+        pass
+
+    def search_embeddings(
+        query_embedding: list[float], top_k: int | None = None
+    ) -> list[tuple[int, float]]:
+        return []

codec/schema_pb2.py

@@ -0,0 +1,61 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: codec/schema.proto
+# Protobuf Python Version: 6.33.1
+"""Generated protocol buffer code."""
+
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import struct_pb2 as _struct_pb2
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC, 6, 33, 1, "", "codec/schema.proto"
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
+    b'\n\x12\x63odec/schema.proto\x12\x0cgraphcode.v1\x1a\x1cgoogle/protobuf/struct.proto"f\n\x0eGraphCodeIndex\x12!\n\x05nodes\x18\x01 \x03(\x0b\x32\x12.graphcode.v1.Node\x12\x31\n\rrelationships\x18\x02 \x03(\x0b\x32\x1a.graphcode.v1.Relationship"\x93\x04\n\x04Node\x12(\n\x07project\x18\x01 \x01(\x0b\x32\x15.graphcode.v1.ProjectH\x00\x12(\n\x07package\x18\x02 \x01(\x0b\x32\x15.graphcode.v1.PackageH\x00\x12&\n\x06\x66older\x18\x03 \x01(\x0b\x32\x14.graphcode.v1.FolderH\x00\x12&\n\x06module\x18\x04 \x01(\x0b\x32\x14.graphcode.v1.ModuleH\x00\x12)\n\nclass_node\x18\x05 \x01(\x0b\x32\x13.graphcode.v1.ClassH\x00\x12*\n\x08\x66unction\x18\x06 \x01(\x0b\x32\x16.graphcode.v1.FunctionH\x00\x12&\n\x06method\x18\x07 \x01(\x0b\x32\x14.graphcode.v1.MethodH\x00\x12"\n\x04\x66ile\x18\x08 \x01(\x0b\x32\x12.graphcode.v1.FileH\x00\x12\x39\n\x10\x65xternal_package\x18\t \x01(\x0b\x32\x1d.graphcode.v1.ExternalPackageH\x00\x12\x43\n\x15module_implementation\x18\n \x01(\x0b\x32".graphcode.v1.ModuleImplementationH\x00\x12\x39\n\x10module_interface\x18\x0b \x01(\x0b\x32\x1d.graphcode.v1.ModuleInterfaceH\x00\x42\t\n\x07payload"\xe9\x03\n\x0cRelationship\x12\x39\n\x04type\x18\x01 \x01(\x0e\x32+.graphcode.v1.Relationship.RelationshipType\x12\x11\n\tsource_id\x18\x02 \x01(\t\x12\x11\n\ttarget_id\x18\x03 \x01(\t\x12+\n\nproperties\x18\x04 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x14\n\x0csource_label\x18\x05 \x01(\t\x12\x14\n\x0ctarget_label\x18\x06 \x01(\t"\x9e\x02\n\x10RelationshipType\x12!\n\x1dRELATIONSHIP_TYPE_UNSPECIFIED\x10\x00\x12\x14\n\x10\x43ONTAINS_PACKAGE\x10\x01\x12\x13\n\x0f\x43ONTAINS_FOLDER\x10\x02\x12\x11\n\rCONTAINS_FILE\x10\x03\x12\x13\n\x0f\x43ONTAINS_MODULE\x10\x04\x12\x0b\n\x07\x44\x45\x46INES\x10\x05\x12\x12\n\x0e\x44\x45\x46INES_METHOD\x10\x06\x12\x0b\n\x07IMPORTS\x10\x07\x12\x0c\n\x08INHERITS\x10\x08\x12\r\n\tOVERRIDES\x10\t\x12\t\n\x05\x43\x41LLS\x10\n\x12\x17\n\x13\x44\x45PENDS_ON_EXTERNAL\x10\x0b\x12\x15\n\x11IMPLEMENTS_MODULE\x10\x0c\x12\x0e\n\nIMPLEMENTS\x10\r"\x17\n\x07Project\x12\x0c\n\x04name\x18\x01 \x01(\t"=\n\x07Package\x12\x16\n\x0equalified_name\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0c\n\x04path\x18\x03 \x01(\t"$\n\x06\x46older\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t"5\n\x04\x46ile\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x11\n\textension\x18\x03 \x01(\t"<\n\x06Module\x12\x16\n\x0equalified_name\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0c\n\x04path\x18\x03 \x01(\t"e\n\x14ModuleImplementation\x12\x16\n\x0equalified_name\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0c\n\x04path\x18\x03 \x01(\t\x12\x19\n\x11implements_module\x18\x04 \x01(\t"E\n\x0fModuleInterface\x12\x16\n\x0equalified_name\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0c\n\x04path\x18\x03 \x01(\t"\x1f\n\x0f\x45xternalPackage\x12\x0c\n\x04name\x18\x01 \x01(\t"\x92\x01\n\x08\x46unction\x12\x16\n\x0equalified_name\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x11\n\tdocstring\x18\x03 \x01(\t\x12\x12\n\nstart_line\x18\x04 \x01(\x05\x12\x10\n\x08\x65nd_line\x18\x05 \x01(\x05\x12\x12\n\ndecorators\x18\x06 \x03(\t\x12\x13\n\x0bis_exported\x18\x07 \x01(\x08"{\n\x06Method\x12\x16\n\x0equalified_name\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x11\n\tdocstring\x18\x03 \x01(\t\x12\x12\n\nstart_line\x18\x04 \x01(\x05\x12\x10\n\x08\x65nd_line\x18\x05 \x01(\x05\x12\x12\n\ndecorators\x18\x06 \x03(\t"\x8f\x01\n\x05\x43lass\x12\x16\n\x0equalified_name\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x11\n\tdocstring\x18\x03 \x01(\t\x12\x12\n\nstart_line\x18\x04 \x01(\x05\x12\x10\n\x08\x65nd_line\x18\x05 \x01(\x05\x12\x12\n\ndecorators\x18\x06 \x03(\t\x12\x13\n\x0bis_exported\x18\x07 \x01(\x08\x62\x06proto3'
+)
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "codec.schema_pb2", _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+    DESCRIPTOR._loaded_options = None
+    _globals["_GRAPHCODEINDEX"]._serialized_start = 66
+    _globals["_GRAPHCODEINDEX"]._serialized_end = 168
+    _globals["_NODE"]._serialized_start = 171
+    _globals["_NODE"]._serialized_end = 702
+    _globals["_RELATIONSHIP"]._serialized_start = 705
+    _globals["_RELATIONSHIP"]._serialized_end = 1194
+    _globals["_RELATIONSHIP_RELATIONSHIPTYPE"]._serialized_start = 908
+    _globals["_RELATIONSHIP_RELATIONSHIPTYPE"]._serialized_end = 1194
+    _globals["_PROJECT"]._serialized_start = 1196
+    _globals["_PROJECT"]._serialized_end = 1219
+    _globals["_PACKAGE"]._serialized_start = 1221
+    _globals["_PACKAGE"]._serialized_end = 1282
+    _globals["_FOLDER"]._serialized_start = 1284
+    _globals["_FOLDER"]._serialized_end = 1320
+    _globals["_FILE"]._serialized_start = 1322
+    _globals["_FILE"]._serialized_end = 1375
+    _globals["_MODULE"]._serialized_start = 1377
+    _globals["_MODULE"]._serialized_end = 1437
+    _globals["_MODULEIMPLEMENTATION"]._serialized_start = 1439
+    _globals["_MODULEIMPLEMENTATION"]._serialized_end = 1540
+    _globals["_MODULEINTERFACE"]._serialized_start = 1542
+    _globals["_MODULEINTERFACE"]._serialized_end = 1611
+    _globals["_EXTERNALPACKAGE"]._serialized_start = 1613
+    _globals["_EXTERNALPACKAGE"]._serialized_end = 1644
+    _globals["_FUNCTION"]._serialized_start = 1647
+    _globals["_FUNCTION"]._serialized_end = 1793
+    _globals["_METHOD"]._serialized_start = 1795
+    _globals["_METHOD"]._serialized_end = 1918
+    _globals["_CLASS"]._serialized_start = 1921
+    _globals["_CLASS"]._serialized_end = 2064
+# @@protoc_insertion_point(module_scope)