o
    *j'                     @   s   d dl Z d dlmZmZmZmZ d dlZd dlmZ d dl	m
Z
 d dlmZ d dlmZmZ d dlmZ d dlmZ d d	lmZ d d
lmZ d dlmZmZ dZdZdZdZdZdgZ ej!ej"ej"dG dd deZ#dS )    N)AnyDictOptionalUnion)	Pipelines)Model)T5Chat)
OutputKeysTokenGeneratorOutput)Pipeline)	PIPELINES)Preprocessor)"remove_space_between_chinese_chars)	ModelFileTasksu~   假设我和你正在进行对话，请你给我得体、准确、友好的回复。以下是我们的对话内容。{context}u   假设我和你正在进行对话，请你给我得体、准确、友好的回复。以下是我们的对话内容。{context}#以下是在此之前我们的对话内容，可作为回复时的参考。{history}u   假设我和你正在进行对话，请你给我得体、准确、友好的回复。以下是我们的对话内容。{context}#以下是和对话相关的知识，请你参考该知识进行回复。{knowledge}u  假设我和你正在进行对话，请你给我得体、准确、友好的回复。以下是我们的对话内容。{context}#假设以下是你对我所了解的信息，请你参考该信息并避免你的回复和该信息矛盾，信息如下：{user_profile}u   假设我和你正在进行对话，请你给我得体、准确、友好的回复。以下是我们的对话内容。{context}#假设以下是你的人物设定，请你参考该信息并避免你的回复和该信息矛盾，信息如下：{bot_profile}FidDialoguePipeline)module_namec                	       s   e Zd Z				ddeeef dee dedef fdd	Zd
d Z	de
eef fddZde
eef de
eef fddZdd Zdd Zdede
eef fddZ  ZS )r   NgpuTmodelpreprocessorconfig_filedevicec           	         s   t  jd
|||||d| t| jt| _|du r&tj| jjfi || _	| jsFg }t
ddD ]}|d| d q0| j	jjd|i t| jtsTJ dtj | j| j| _| j  d	| _dS )u  Use `model` and `preprocessor` to create a fid-dialogue pipeline for prediction.

        Args:
            model (str or Model): Supply either a local model dir which supported the text generation task,
            or a model id from the model hub, or a torch model instance.
            preprocessor (Preprocessor): An optional preprocessor instance, please make sure the preprocessor fits for
            the model if supplied.
            kwargs (dict, `optional`):
                Extra kwargs passed into the preprocessor's constructor.
            Examples:
                >>> from modelscope.pipelines import pipeline
                >>> from modelscope.utils.constant import Tasks
                >>> pipeline_ins = pipeline(Tasks.fid_dialogue, model='damo/plug-dialogue', model_revision='v1.0.1')
                >>> input = {
                >>>    "history": "你好[SEP]你好，我是小达，很高兴认识你！[SEP]李白是谁",
                >>>    "bot_profile": "我是小达;我是女生;我是单身;我今年21岁;我生日是2001年11月11日",
                >>>    "knowledge": "唐代诗人李白（701年—762年12月）,字太白,号青莲居士,又号“谪仙人”[SEP]李白（公元701年—公元762年），字太白",
                >>>    "user_profile": "你是小明"
                >>> }
                >>> result = pipeline_ins(input)
                >>> print(result)
        )r   r   r   r   auto_collateN   d   z[unused]Zadditional_special_tokensz,please check whether model config exists in [SEP] )super__init__
isinstancer   r   is_t5r   Zfrom_pretrainedZ	model_dirpreprocessor_tokenizerrangeappendZnlp_tokenizer	tokenizerZadd_special_tokensr   r   ZCONFIGURATIONtor   evalSEP)	selfr   r   r   r   r   kwargsZunused_listi	__class__r   o/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/modelscope/pipelines/nlp/fid_dialogue_pipeline.pyr   #   s:   



zFidDialoguePipeline.__init__c                 K   s.   | di }| di }| di }|||fS )Npreprocess_paramsforward_paramspostprocess_params)get)r)   Zpipeline_parametersr/   r0   r1   r   r   r.   _sanitize_parameters[   s   
z(FidDialoguePipeline._sanitize_parametersinputsc                 K   s@   t   | jj|fi |W  d    S 1 sw   Y  d S )N)torchZno_gradr   generate)r)   r4   r0   r   r   r.   forwarda   s   
$zFidDialoguePipeline.forwardreturnc                 K   s4  d}d}d|v r| d}d|v r| d}d|v r|d nd}t|dkr*tdd	|v r2|d	 nd}d
|v r<|d
 nd}d|v rF|d nd}|| j}|| d  }	| |	}	|d |  }| |}|| j}g }
|rt|dkr|
tj	|	|d |rt|dkr|D ]}|
t
j	|	|d q|rt|dkr|
tj	|	|d |rt|dkr|
tj	|	|d |
s|
tj	|	d tt|
D ]*}| jr|
| dddddd|
|< q|
| dddddd|
|< q| jd|
idd|ddd d| j}d|tj| ji}|S )Ni,     max_encoder_lengthcontext_turnhistory r   zhistory is necessary!	knowledgeuser_profilebot_profile)contextr<   )rA   r>   )rA   r?   )rA   r@   )rA   
u   ▁<extra_id_22>	u   ▁<extra_id_33>  u   ▁<extra_id_23>
[unused22]
[unused33]
[unused23]Zsrc_txtTpt)paddingZ
truncation
max_lengthZreturn_tensors	input_ids)poplen	Exceptionsplitr(   process_contextprocess_historyr$   history_templateformatknowledge_templateuser_profile_templatebot_profile_templatecontext_templater#   r!   replacer"   Z	unsqueezer&   r   r5   Zint64)r)   r4   r/   r:   r;   r<   r>   r?   r@   rA   Zmodel_inputZknowr+   rK   Z
input_dictr   r   r.   
preprocesse   s   




zFidDialoguePipeline.preprocessc                 C   s|   d}t t|d ddD ],}t|| dkr&|| d dvr&|| d ||< |d ||  ||< |dkr6dnd}qd	|S )
N   我r   r   Y   !"#$%&'()*+,-./:;<=>?@[\]^_`{|}~、。，？！；：“”（）【】《》〈〉……   。   ：   你r=   r#   rM   join)r)   Zcontext_listsubjectr+   r   r   r.   rP         
z#FidDialoguePipeline.process_contextc                 C   s|   d}t t|d ddD ],}t|| dkr&|| d dvr&|| d ||< |d ||  ||< |dkr6dnd}qd	|S )
Nr_   r   r[   r   r\   r]   r^   rZ   r=   r`   )r)   Zhistory_listrb   r+   r   r   r.   rQ      rc   z#FidDialoguePipeline.process_historyc                 K   sx   |j    }| jj|d | jd}ddddddddddd
}| D ]
\}}|||}q%| js7t	|}t
j|iS )Nr   )Zskip_special_tokensrB   rC   rD   r=   )
z<extra_id_22>z<extra_id_33>z<extra_id_23>rE   rF   rG   r   z[CLS]z[PAD]z[UNK])	sequencesdetachcputolistr"   decoder!   itemsrX   r   r	   ZTEXT)r)   r4   r1   Z
hypothesesresponseZtoken_mappingstr   r   r.   postprocess   s(   

zFidDialoguePipeline.postprocess)NNr   T)__name__
__module____qualname__r   r   strr   r   r   r3   r   r   r7   rY   rP   rQ   r
   rm   __classcell__r   r   r,   r.   r      s0    
8

F


)$retypingr   r   r   r   r5   Zmodelscope.metainfor   Zmodelscope.models.baser   Z,modelscope.models.nlp.fid_T5.text_generationr   Zmodelscope.outputsr	   r
   Zmodelscope.pipelines.baser   Zmodelscope.pipelines.builderr   Zmodelscope.preprocessorsr   Zmodelscope.utils.chinese_utilsr   Zmodelscope.utils.constantr   r   rW   rR   rT   rU   rV   __all__Zregister_moduleZfid_dialoguer   r   r   r   r.   <module>   s,   