mirror of
https://github.com/infiniflow/ragflow.git
synced 2025-12-08 20:42:30 +08:00
add support for LocalAI (#1608)
### What problem does this PR solve? #762 ### Type of change - [x] New Feature (non-breaking change which adds functionality) --------- Co-authored-by: Zhedong Cen <cenzhedong2@126.com>
This commit is contained in:
@ -111,6 +111,24 @@ class OpenAIEmbed(Base):
|
||||
return np.array(res.data[0].embedding), res.usage.total_tokens
|
||||
|
||||
|
||||
class LocalAIEmbed(Base):
|
||||
def __init__(self, key, model_name, base_url):
|
||||
self.base_url = base_url + "/embeddings"
|
||||
self.headers = {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
self.model_name = model_name.split("___")[0]
|
||||
|
||||
def encode(self, texts: list, batch_size=None):
|
||||
data = {"model": self.model_name, "input": texts, "encoding_type": "float"}
|
||||
res = requests.post(self.base_url, headers=self.headers, json=data).json()
|
||||
|
||||
return np.array([d["embedding"] for d in res["data"]]), 1024
|
||||
|
||||
def encode_queries(self, text):
|
||||
embds, cnt = self.encode([text])
|
||||
return np.array(embds[0]), cnt
|
||||
|
||||
class AzureEmbed(OpenAIEmbed):
|
||||
def __init__(self, key, model_name, **kwargs):
|
||||
self.client = AzureOpenAI(api_key=key, azure_endpoint=kwargs["base_url"], api_version="2024-02-01")
|
||||
@ -443,4 +461,4 @@ class GeminiEmbed(Base):
|
||||
task_type="retrieval_document",
|
||||
title="Embedding of single string")
|
||||
token_count = num_tokens_from_string(text)
|
||||
return np.array(result['embedding']),token_count
|
||||
return np.array(result['embedding']),token_count
|
||||
|
||||
Reference in New Issue
Block a user