mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-23 00:43:18 +00:00
59 lines
2.2 KiB
Python
59 lines
2.2 KiB
Python
import os
|
|
from typing import Literal, List
|
|
|
|
from dotenv import load_dotenv
|
|
from openai import OpenAI
|
|
from pydantic import Field, PrivateAttr, model_validator
|
|
|
|
from experiencemaker.embedding_model import EMBEDDING_MODEL_REGISTRY
|
|
from experiencemaker.embedding_model.base_embedding_model import BaseEmbeddingModel
|
|
|
|
|
|
@EMBEDDING_MODEL_REGISTRY.register("openai_compatible")
|
|
class OpenAICompatibleEmbeddingModel(BaseEmbeddingModel):
|
|
api_key: str = Field(default_factory=lambda: os.getenv("EMBEDDING_API_KEY"), description="api key")
|
|
base_url: str = Field(default_factory=lambda: os.getenv("EMBEDDING_BASE_URL"), description="base url")
|
|
model_name: str = Field(default="", description="model name")
|
|
dimensions: int = Field(default=1024, description="dimensions")
|
|
encoding_format: Literal["float", "base64"] = Field(default="float", description="encoding_format")
|
|
_client: OpenAI = PrivateAttr()
|
|
|
|
@model_validator(mode="after")
|
|
def init_client(self):
|
|
self._client = OpenAI(api_key=self.api_key, base_url=self.base_url)
|
|
return self
|
|
|
|
def _get_embeddings(self, input_text: str | List[str]):
|
|
completion = self._client.embeddings.create(
|
|
model=self.model_name,
|
|
input=input_text,
|
|
dimensions=self.dimensions,
|
|
encoding_format=self.encoding_format
|
|
)
|
|
|
|
if isinstance(input_text, str):
|
|
return completion.data[0].embedding
|
|
|
|
elif isinstance(input_text, list):
|
|
result_emb = [[] for _ in range(len(input_text))]
|
|
for emb in completion.data:
|
|
result_emb[emb.index] = emb.embedding
|
|
return result_emb
|
|
|
|
else:
|
|
raise RuntimeError(f"unsupported type={type(input_text)}")
|
|
|
|
|
|
def main():
|
|
load_dotenv()
|
|
model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4")
|
|
res1 = model.get_embeddings(
|
|
"The clothes are of good quality and look good, definitely worth the wait. I love them.")
|
|
res2 = model.get_embeddings(["aa", "bb"])
|
|
print(res1)
|
|
print(res2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
# launch with: python -m experiencemaker.model.openai_compatible_embedding_model
|