Spaces:
Sleeping
Sleeping
File size: 3,401 Bytes
308de05 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
import logging
from ntr_fileparser import ParsedDocument
from ntr_text_fragmentation import (
ChunkingStrategy,
LinkerEntity,
register_chunking_strategy,
register_entity,
DocumentAsEntity,
Chunk,
)
from components.llm.common import LlmPredictParams
from components.llm.deepinfra_api import DeepInfraApi
from components.llm.prompts import PROMPT_APPENDICES
from components.services.llm_config import LLMConfigService
logger = logging.getLogger(__name__)
APPENDICES_CHUNKER = 'appendices'
@register_entity
class Appendix(Chunk):
"""Сущность для хранения приложений"""
@register_chunking_strategy(APPENDICES_CHUNKER)
class AppendicesProcessor(ChunkingStrategy):
def __init__(
self,
llm_api: DeepInfraApi,
llm_config_service: LLMConfigService,
):
self.prompt = PROMPT_APPENDICES
self.llm_api = llm_api
p = llm_config_service.get_default()
self.llm_params = LlmPredictParams(
temperature=p.temperature,
top_p=p.top_p,
min_p=p.min_p,
seed=p.seed,
frequency_penalty=p.frequency_penalty,
presence_penalty=p.presence_penalty,
n_predict=p.n_predict,
)
def chunk(
self, document: ParsedDocument, doc_entity: DocumentAsEntity
) -> list[LinkerEntity]:
raise NotImplementedError(
f"{self.__class__.__name__} поддерживает только асинхронный вызов. "
"Используйте метод extract_async или другую стратегию."
)
async def chunk_async(
self, document: ParsedDocument, doc_entity: DocumentAsEntity
) -> list[LinkerEntity]:
text = ""
text += document.name + "\n"
text += "\n".join([p.text for p in document.paragraphs])
text += "\n".join([t.to_string() for t in document.tables])
prompt = self._format_prompt(text)
response = await self.llm_api.predict(prompt=prompt, system_prompt=None)
processed = self._postprocess_llm_response(response)
if processed is None:
return []
entity = Appendix(
text=processed,
in_search_text=processed,
number_in_relation=0,
groupper=APPENDICES_CHUNKER,
)
entity.owner_id = doc_entity.id
return [entity]
def _format_prompt(self, text: str) -> str:
return self.prompt.format(replace_me=text)
def _postprocess_llm_response(self, response: str | None) -> str | None:
if response is None:
return None
# Найти начало и конец текста в квадратных скобках
start = response.find('[')
end = response.find(']')
# Проверка, что найдена только одна пара скобок
if start == -1 or end == -1 or start >= end:
logger.warning(f"Некорректный формат ответа LLM: {response}")
return None
# Извлечь текст внутри скобок
extracted_text = response[start + 1 : end]
if extracted_text == '%%':
logging.info(f'Приложение признано бесполезным')
return None
return extracted_text
|