49 lines
1.3 KiB
Python
49 lines
1.3 KiB
Python
# coding=utf-8
|
||
"""
|
||
@project: MaxKB
|
||
@Author：虎
|
||
@file： embedding.py
|
||
@date：2024/7/12 15:02
|
||
@desc:
|
||
"""
|
||
from typing import Dict, List
|
||
|
||
from langchain_community.embeddings import OllamaEmbeddings
|
||
|
||
from models_provider.base_model_provider import MaxKBBaseModel
|
||
|
||
|
||
class OllamaEmbedding(MaxKBBaseModel, OllamaEmbeddings):
|
||
@staticmethod
|
||
def new_instance(model_type, model_name, model_credential: Dict[str, object], **model_kwargs):
|
||
return OllamaEmbedding(
|
||
model=model_name,
|
||
base_url=model_credential.get('api_base'),
|
||
)
|
||
|
||
def embed_documents(self, texts: List[str]) -> List[List[float]]:
|
||
"""Embed documents using an Ollama deployed embedding model.
|
||
|
||
Args:
|
||
texts: The list of texts to embed.
|
||
|
||
Returns:
|
||
List of embeddings, one for each text.
|
||
"""
|
||
instruction_pairs = [f"{text}" for text in texts]
|
||
embeddings = self._embed(instruction_pairs)
|
||
return embeddings
|
||
|
||
def embed_query(self, text: str) -> List[float]:
|
||
"""Embed a query using a Ollama deployed embedding model.
|
||
|
||
Args:
|
||
text: The text to embed.
|
||
|
||
Returns:
|
||
Embeddings for the text.
|
||
"""
|
||
instruction_pair = f"{text}"
|
||
embedding = self._embed([instruction_pair])[0]
|
||
return embedding
|