Spaces:
Sleeping
Sleeping
File size: 1,201 Bytes
ed83bfd bc0a1d6 ed83bfd |
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 |
import os
from dotenv import load_dotenv
from typing import Optional
from langchain_openai.embeddings import OpenAIEmbeddings
load_dotenv()
class MyOpenAIEmbeddings:
@classmethod
def from_model(
cls,
model: str = 'small',
*,
dimensions: Optional[int] = None,
max_retries: int = 1,
**kwargs
)-> OpenAIEmbeddings:
if model in ['text-embedding-3-small', 'TEXT-EMBEDDING-3-SMALL', 'small', 'SMALL']:
model = 'text-embedding-3-small'
dimensions = 1536 if dimensions is None else dimensions
elif model in ['text-embedding-3-large', 'TEXT-EMBEDDING-3-LARGE', 'large', 'LARGE']:
model = 'text-embedding-3-large'
dimensions = 3072 if dimensions is None else dimensions
else:
raise ValueError(f"Model {model} is currently not supported. Supported models are: ['text-embedding-3-small', 'text-embedding-3-large']")
return OpenAIEmbeddings(
openai_api_key=os.getenv("OPENAI_API_KEY"),
model=model,
dimensions=dimensions,
max_retries=max_retries,
**kwargs
)
|