| from __future__ import annotations |
|
|
| from typing import Iterable |
|
|
| from app.core.config import settings |
|
|
| try: |
| import torch |
| from transformers import AutoModel, AutoTokenizer |
| except Exception: |
| torch = None |
| AutoModel = None |
| AutoTokenizer = None |
|
|
|
|
| class CodeEmbeddingService: |
| def __init__(self) -> None: |
| self._dim = settings.embedding_dim |
| self._ready = False |
| self._tokenizer = None |
| self._model = None |
|
|
| if AutoModel is not None and AutoTokenizer is not None: |
| try: |
| self._tokenizer = AutoTokenizer.from_pretrained(settings.codebert_model) |
| self._model = AutoModel.from_pretrained(settings.codebert_model) |
| self._model.eval() |
| self._ready = True |
| except Exception: |
| self._ready = False |
|
|
| @property |
| def ready(self) -> bool: |
| return self._ready |
|
|
| @property |
| def embedding_dim(self) -> int: |
| return self._dim |
|
|
| def embed_repository_signals(self, snippets: Iterable[str]) -> list[float]: |
| text = "\n".join(snippets).strip() |
| if not text: |
| return [0.0] * self._dim |
|
|
| if self._ready and torch is not None and self._tokenizer is not None and self._model is not None: |
| tokens = self._tokenizer(text, return_tensors="pt", truncation=True, max_length=512) |
| with torch.no_grad(): |
| outputs = self._model(**tokens) |
| cls = outputs.last_hidden_state[:, 0, :].squeeze(0) |
| return cls.tolist() |
|
|
| vector = [0.0] * self._dim |
| for i, ch in enumerate(text): |
| vector[i % self._dim] += (ord(ch) % 31) / 31.0 |
| norm = sum(abs(v) for v in vector) or 1.0 |
| return [v / norm for v in vector] |
|
|
