File size: 1,445 Bytes
81598c5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
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
"""

Azure AI Foundry embedding provider.

Uses text-embedding-ada-002 (or whatever deployment is configured).

"""

from openai import AzureOpenAI
from openmark.embeddings.base import EmbeddingProvider
from openmark import config


class AzureEmbedder(EmbeddingProvider):
    def __init__(self):
        self._client = AzureOpenAI(
            azure_endpoint=config.AZURE_ENDPOINT,
            api_key=config.AZURE_API_KEY,
            api_version=config.AZURE_API_VERSION,
        )
        self._deployment = config.AZURE_DEPLOYMENT_EMBED
        print(f"Azure embedder ready — deployment: {self._deployment}")

    def _embed(self, texts: list[str]) -> list[list[float]]:
        response = self._client.embeddings.create(
            input=texts,
            model=self._deployment,
        )
        return [item.embedding for item in response.data]

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        results = []
        batch_size = 100
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]
            results.extend(self._embed(batch))
            print(f"  Azure embedded {min(i + batch_size, len(texts))}/{len(texts)}")
        return results

    def embed_query(self, text: str) -> list[float]:
        return self._embed([text])[0]

    @property
    def dimension(self) -> int:
        return 1536  # ada-002 dimension