wuenlp_tools.utils.embedding 

 1from dataclasses import dataclass
 2from typing import List, Optional
 3
 4from openai import OpenAI, NOT_GIVEN
 5from loguru import logger
 6
 7from wuenlp_tools.keys import OPENAI_API_KEY
 8
 9
10@dataclass
11class EmbeddingModel:
12    name: str
13    model: str
14    dimensions: Optional[int] = None
15    max_tokens: int = 0
16    openai: bool = True
17    openrouter: bool = False
18
19
20OpenAIEmbeddingSmall = EmbeddingModel(name="text-embedding-3-small", model="text-embedding-3-small", dimensions=512,
21                                      max_tokens=8192)
22OpenAIEmbeddingLarge = EmbeddingModel(name="text-embedding-3-large", model="text-embedding-3-large")
23
24
25class TextEmbedder:
26    openai: OpenAI
27    model: EmbeddingModel
28
29    def __init__(self, model: EmbeddingModel):
30        if not model.openai:
31            raise NotImplementedError("Non-OpenAI embedding models are currently not supported")
32        self.openai = OpenAI(api_key=OPENAI_API_KEY)
33        self.model = model
34
35    def __call__(self, text: str) -> List[float]:
36        dimensions = self.model.dimensions if self.model.dimensions is not None else NOT_GIVEN
37        if len(text) > self.model.max_tokens:
38            logger.info(f"Text too long: {len(text)} > {self.model.max_tokens}. Cutting.")
39            text = text[:self.model.max_tokens]
40        response = self.openai.embeddings.create(input=text, model=self.model.model, dimensions=dimensions)
41
42        return response.data[0].embedding
43
44
45if __name__ == '__main__':
46    embedder = TextEmbedder(OpenAIEmbeddingSmall)
47
48    print(embedder("hello world"))
@dataclass
class EmbeddingModel: 
11@dataclass
12class EmbeddingModel:
13    name: str
14    model: str
15    dimensions: Optional[int] = None
16    max_tokens: int = 0
17    openai: bool = True
18    openrouter: bool = False
EmbeddingModel( name: str, model: str, dimensions: Optional[int] = None, max_tokens: int = 0, openai: bool = True, openrouter: bool = False)
name: str
model: str
dimensions: Optional[int] = None
max_tokens: int = 0

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.__int__(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer literal.

>>> int('0b100', base=0)
4
openai: bool = True

bool(x) -> bool

Returns True when the argument x is true, False otherwise. The builtins True and False are the only two instances of the class bool. The class bool is a subclass of the class int, and cannot be subclassed.

openrouter: bool = False

bool(x) -> bool

Returns True when the argument x is true, False otherwise. The builtins True and False are the only two instances of the class bool. The class bool is a subclass of the class int, and cannot be subclassed.

OpenAIEmbeddingSmall = EmbeddingModel(name='text-embedding-3-small', model='text-embedding-3-small', dimensions=512, max_tokens=8192, openai=True, openrouter=False)

EmbeddingModel(name: str, model: str, dimensions: Optional[int] = None, max_tokens: int = 0, openai: bool = True, openrouter: bool = False)

OpenAIEmbeddingLarge = EmbeddingModel(name='text-embedding-3-large', model='text-embedding-3-large', dimensions=None, max_tokens=0, openai=True, openrouter=False)

EmbeddingModel(name: str, model: str, dimensions: Optional[int] = None, max_tokens: int = 0, openai: bool = True, openrouter: bool = False)

class TextEmbedder: 
26class TextEmbedder:
27    openai: OpenAI
28    model: EmbeddingModel
29
30    def __init__(self, model: EmbeddingModel):
31        if not model.openai:
32            raise NotImplementedError("Non-OpenAI embedding models are currently not supported")
33        self.openai = OpenAI(api_key=OPENAI_API_KEY)
34        self.model = model
35
36    def __call__(self, text: str) -> List[float]:
37        dimensions = self.model.dimensions if self.model.dimensions is not None else NOT_GIVEN
38        if len(text) > self.model.max_tokens:
39            logger.info(f"Text too long: {len(text)} > {self.model.max_tokens}. Cutting.")
40            text = text[:self.model.max_tokens]
41        response = self.openai.embeddings.create(input=text, model=self.model.model, dimensions=dimensions)
42
43        return response.data[0].embedding
TextEmbedder(model: EmbeddingModel) 
30    def __init__(self, model: EmbeddingModel):
31        if not model.openai:
32            raise NotImplementedError("Non-OpenAI embedding models are currently not supported")
33        self.openai = OpenAI(api_key=OPENAI_API_KEY)
34        self.model = model
openai: openai.OpenAI