Ascend/gwdoc
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

优化models_api

Browse Source
This commit is contained in:
Defeng 2026-07-17 15:19:55 +08:00
parent 93ced6ceb0
commit cd70022d67
1 changed files with 0 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
modelsAPI/model_api.py
View File

@ -5,7 +5,6 @@ import os
import numpy as np import numpy as np
from openai import OpenAI from openai import OpenAI
from config import LLM_CONFIG, EMBEDDING_CONFIG from config import LLM_CONFIG, EMBEDDING_CONFIG
from neo4j_graphrag.embeddings.base import Embedder
from openai import AsyncOpenAI from openai import AsyncOpenAI
@ -353,24 +352,7 @@ class OpenaiAPI:
return 0.5 * jaccard + 0.5 * len_penalty return 0.5 * jaccard + 0.5 * len_penalty
class LocalBgeM3Embeddings(Embedder):
def __init__(self, base_url: str, api_key: str, model_name: str = "bge-m3"):
self.base_url = base_url
self.api_key = api_key
self.model_name = model_name
def embed_query(self, text: str) -> List[float]:
return self.embed_documents([text])[0]
def embed_documents(self, texts: List[str]) -> List[List[float]]:
url = self.base_url
with httpx.Client(timeout=60) as client:
response = client.post(
url,
json={"model": self.model_name, "input": texts},
headers={"Authorization": f"Bearer {self.api_key}"},
)
return [item["embedding"] for item in response.json()["data"]]
if __name__ == "__main__": if __name__ == "__main__":