优化models_api

This commit is contained in:
Defeng 2026-07-17 15:19:55 +08:00
parent 93ced6ceb0
commit cd70022d67

View File

@ -5,7 +5,6 @@ import os
import numpy as np
from openai import OpenAI
from config import LLM_CONFIG, EMBEDDING_CONFIG
from neo4j_graphrag.embeddings.base import Embedder
from openai import AsyncOpenAI
@ -353,24 +352,7 @@ class OpenaiAPI:
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__":