77 lines
2.9 KiB
Python
77 lines
2.9 KiB
Python
|
# coding=utf-8
|
|||
|
|
"""
|
|||
|
|
@project: MaxKB
|
|||
|
|
@Author：虎
|
|||
|
|
@file： model_apply_serializers.py
|
|||
|
|
@date：2024/8/20 20:39
|
|||
|
|
@desc:
|
|||
|
|
"""
|
|||
|
|
from django.db import connection
|
|||
|
|
from django.db.models import QuerySet
|
|||
|
|
from langchain_core.documents import Document
|
|||
|
|
from rest_framework import serializers
|
|||
|
|
|
|||
|
|
from common.config.embedding_config import ModelManage
|
|||
|
|
from django.utils.translation import gettext_lazy as _
|
|||
|
|
|
|||
|
|
from models_provider.models import Model
|
|||
|
|
from models_provider.tools import get_model
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_embedding_model(model_id):
|
|||
|
|
model = QuerySet(Model).filter(id=model_id).first()
|
|||
|
|
# 手动关闭数据库连接
|
|||
|
|
connection.close()
|
|||
|
|
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,
|
|||
|
|
label=_('vector text')),
|
|||
|
|
label=_('vector text list')),
|
|||
|
|
|
|||
|
|
|
|||
|
|
class EmbedQuery(serializers.Serializer):
|
|||
|
|
text = serializers.CharField(required=True, label=_('vector text'))
|
|||
|
|
|
|||
|
|
|
|||
|
|
class CompressDocument(serializers.Serializer):
|
|||
|
|
page_content = serializers.CharField(required=True, label=_('text'))
|
|||
|
|
metadata = serializers.DictField(required=False, label=_('metadata'))
|
|||
|
|
|
|||
|
|
|
|||
|
|
class CompressDocuments(serializers.Serializer):
|
|||
|
|
documents = CompressDocument(required=True, many=True)
|
|||
|
|
query = serializers.CharField(required=True, label=_('query'))
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ModelApplySerializers(serializers.Serializer):
|
|||
|
|
model_id = serializers.UUIDField(required=True, label=_('model 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'))
|
|||
|
|
|
|||
|
|
def compress_documents(self, instance, with_valid=True):
|
|||
|
|
if with_valid:
|
|||
|
|
self.is_valid(raise_exception=True)
|
|||
|
|
CompressDocuments(data=instance).is_valid(raise_exception=True)
|
|||
|
|
model = get_embedding_model(self.data.get('model_id'))
|
|||
|
|
return [{'page_content': d.page_content, 'metadata': d.metadata} for d in model.compress_documents(
|
|||
|
|
[Document(page_content=document.get('page_content'), metadata=document.get('metadata')) for document in
|
|||
|
|
instance.get('documents')], instance.get('query'))]
|