54 lines
2.0 KiB
Python
54 lines
2.0 KiB
Python
|
# coding=utf-8
|
|||
|
|
"""
|
|||
|
|
@project: MaxKB
|
|||
|
|
@Author：虎
|
|||
|
|
@file： model_apply_serializers.py
|
|||
|
|
@date：2024/8/20 20:39
|
|||
|
|
@desc:
|
|||
|
|
"""
|
|||
|
|
from django.db.models import QuerySet
|
|||
|
|
from rest_framework import serializers
|
|||
|
|
|
|||
|
|
from common.config.embedding_config import ModelManage
|
|||
|
|
from common.util.field_message import ErrMessage
|
|||
|
|
from setting.models import Model
|
|||
|
|
from setting.models_provider import get_model
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_embedding_model(model_id):
|
|||
|
|
model = QuerySet(Model).filter(id=model_id).first()
|
|||
|
|
embedding_model = ModelManage.get_model(model_id,
|
|||
|
|
lambda _id: get_model(model, use_local=True))
|
|||
|
|
return embedding_model
|
|||
|
|
|
|||
|
|
|
|||
|
|
class EmbedDocuments(serializers.Serializer):
|
|||
|
|
texts = serializers.ListField(required=True, child=serializers.CharField(required=True,
|
|||
|
|
error_messages=ErrMessage.char(
|
|||
|
|
"向量文本")),
|
|||
|
|
error_messages=ErrMessage.list("向量文本列表"))
|
|||
|
|
|
|||
|
|
|
|||
|
|
class EmbedQuery(serializers.Serializer):
|
|||
|
|
text = serializers.CharField(required=True, error_messages=ErrMessage.char("向量文本"))
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ModelApplySerializers(serializers.Serializer):
|
|||
|
|
model_id = serializers.UUIDField(required=True, error_messages=ErrMessage.uuid("模型id"))
|
|||
|
|
|
|||
|
|
def embed_documents(self, instance, with_valid=True):
|
|||
|
|
if with_valid:
|
|||
|
|
self.is_valid(raise_exception=True)
|
|||
|
|
EmbedDocuments(data=instance).is_valid(raise_exception=True)
|
|||
|
|
|
|||
|
|
model = get_embedding_model(self.data.get('model_id'))
|
|||
|
|
return model.embed_documents(instance.getlist('texts'))
|
|||
|
|
|
|||
|
|
def embed_query(self, instance, with_valid=True):
|
|||
|
|
if with_valid:
|
|||
|
|
self.is_valid(raise_exception=True)
|
|||
|
|
EmbedQuery(data=instance).is_valid(raise_exception=True)
|
|||
|
|
|
|||
|
|
model = get_embedding_model(self.data.get('model_id'))
|
|||
|
|
return model.embed_query(instance.get('text'))
|