178 lines
6.4 KiB
Python
178 lines
6.4 KiB
Python
from abc import ABC, abstractmethod
|
||
from typing import Any, Dict, List, Optional, Generator
|
||
import traceback
|
||
from collections import deque
|
||
from queue import Queue
|
||
from threading import Thread
|
||
from langchain.callbacks.manager import CallbackManagerForChainRun
|
||
from models.loader import LoaderCheckPoint
|
||
from pydantic import BaseModel
|
||
import torch
|
||
import transformers
|
||
|
||
|
||
class ListenerToken:
|
||
"""
|
||
观测结果
|
||
"""
|
||
|
||
input_ids: torch.LongTensor
|
||
_scores: torch.FloatTensor
|
||
|
||
def __init__(self, input_ids: torch.LongTensor, _scores: torch.FloatTensor):
|
||
self.input_ids = input_ids
|
||
self._scores = _scores
|
||
|
||
|
||
class AnswerResult(BaseModel):
|
||
"""
|
||
消息实体
|
||
"""
|
||
history: List[List[str]] = []
|
||
llm_output: Optional[dict] = None
|
||
|
||
|
||
class AnswerResultStream:
|
||
def __init__(self, callback_func=None):
|
||
self.callback_func = callback_func
|
||
|
||
def __call__(self, answerResult: AnswerResult):
|
||
if self.callback_func is not None:
|
||
self.callback_func(answerResult)
|
||
|
||
|
||
class AnswerResultQueueSentinelTokenListenerQueue(transformers.StoppingCriteria):
|
||
"""
|
||
定义模型stopping_criteria 监听者，在每次响应时将队列数据同步到AnswerResult
|
||
实现此监听器的目的是，不同模型的预测输出可能不是矢量信息，hf框架可以自定义transformers.StoppingCriteria入参来接收每次预测的Tensor和损失函数，
|
||
通过给 StoppingCriteriaList指定模型生成答案时停止的条件。每个 StoppingCriteria 对象表示一个停止条件
|
||
当每轮预测任务开始时，StoppingCriteria都会收到相同的预测结果，最终由下层实现类确认是否结束
|
||
输出值可用于 generatorAnswer generate_with_streaming的自定义参数观测，以实现更加精细的控制
|
||
"""
|
||
|
||
listenerQueue: deque = deque(maxlen=1)
|
||
|
||
def __init__(self):
|
||
transformers.StoppingCriteria.__init__(self)
|
||
|
||
def __call__(self, input_ids: torch.LongTensor, _scores: torch.FloatTensor, **kwargs) -> bool:
|
||
"""
|
||
每次响应时将数据添加到响应队列
|
||
:param input_ids:
|
||
:param _scores:
|
||
:param kwargs:
|
||
:return:
|
||
"""
|
||
self.listenerQueue.append(ListenerToken(input_ids=input_ids, _scores=_scores))
|
||
return False
|
||
|
||
|
||
class Iteratorize:
|
||
"""
|
||
Transforms a function that takes a callback
|
||
into a lazy iterator (generator).
|
||
"""
|
||
|
||
def __init__(self, func, kwargs={}):
|
||
self.mfunc = func
|
||
self.q = Queue()
|
||
self.sentinel = object()
|
||
self.kwargs = kwargs
|
||
self.stop_now = False
|
||
|
||
def _callback(val):
|
||
"""
|
||
模型输出预测结果收集
|
||
通过定义generate_with_callback收集器AnswerResultStream，收集模型预测的AnswerResult响应结果，最终由下层实现类确认是否结束
|
||
结束条件包含如下
|
||
1、模型预测结束、收集器self.q队列收到 self.sentinel标识
|
||
2、在处理迭代器队列消息时返回了break跳出迭代器，触发了StopIteration事件
|
||
3、模型预测出错
|
||
因为当前类是迭代器，所以在for in 中执行了break后 __exit__ 方法会被调用，最终stop_now属性会被更新，然后抛出异常结束预测行为
|
||
迭代器收集的行为如下
|
||
创建Iteratorize迭代对象，
|
||
定义generate_with_callback收集器AnswerResultStream
|
||
启动一个线程异步预测结果来调用上游checkpoint的实现方法_generate_answer
|
||
_generate_answer通过generate_with_callback定义的收集器，收集上游checkpoint包装的AnswerResult消息体
|
||
由于self.q是阻塞模式，每次预测后会被消费后才会执行下次预测
|
||
这时generate_with_callback会被阻塞
|
||
主线程Iteratorize对象的__next__方法调用获取阻塞消息并消费
|
||
1、消息为上游checkpoint包装的AnswerResult消息体，返回下游处理
|
||
2、消息为self.sentinel标识，抛出StopIteration异常
|
||
主线程Iteratorize对象__exit__收到消息，最终stop_now属性会被更新
|
||
异步线程检测stop_now属性被更新，抛出异常结束预测行为
|
||
迭代行为结束
|
||
:param val:
|
||
:return:
|
||
"""
|
||
if self.stop_now:
|
||
raise ValueError
|
||
self.q.put(val)
|
||
|
||
def gen():
|
||
try:
|
||
ret = self.mfunc(callback=_callback, **self.kwargs)
|
||
except ValueError:
|
||
pass
|
||
except:
|
||
traceback.print_exc()
|
||
pass
|
||
|
||
self.q.put(self.sentinel)
|
||
|
||
self.thread = Thread(target=gen)
|
||
self.thread.start()
|
||
|
||
def __iter__(self):
|
||
return self
|
||
|
||
def __next__(self):
|
||
obj = self.q.get(True, None)
|
||
if obj is self.sentinel:
|
||
raise StopIteration
|
||
else:
|
||
return obj
|
||
|
||
def __del__(self):
|
||
"""
|
||
暂无实现
|
||
:return:
|
||
"""
|
||
pass
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
""" break 后会执行 """
|
||
self.stop_now = True
|
||
|
||
|
||
class BaseAnswer(ABC):
|
||
"""上层业务包装器.用于结果生成统一api调用"""
|
||
|
||
@property
|
||
@abstractmethod
|
||
def _check_point(self) -> LoaderCheckPoint:
|
||
"""Return _check_point of llm."""
|
||
def generatorAnswer(self,
|
||
inputs: Dict[str, Any],
|
||
run_manager: Optional[CallbackManagerForChainRun] = None,) -> Generator[Any, str, bool]:
|
||
def generate_with_callback(callback=None, **kwargs):
|
||
kwargs['generate_with_callback'] = AnswerResultStream(callback_func=callback)
|
||
self._generate_answer(**kwargs)
|
||
|
||
def generate_with_streaming(**kwargs):
|
||
return Iteratorize(generate_with_callback, kwargs)
|
||
|
||
with generate_with_streaming(inputs=inputs, run_manager=run_manager) as generator:
|
||
for answerResult in generator:
|
||
yield answerResult
|
||
|
||
@abstractmethod
|
||
def _generate_answer(self,
|
||
inputs: Dict[str, Any],
|
||
run_manager: Optional[CallbackManagerForChainRun] = None,
|
||
generate_with_callback: AnswerResultStream = None) -> None:
|
||
pass
|