a
    bgv5                     @  s   d dl mZ d dlZd dlZd dlmZmZmZm	Z	m
Z
 d dlmZ d dlmZ d dlmZ d dlmZmZmZ d dlmZ d	Zd
ZdZdZeeZeddddG dd deZdS )    )annotationsN)AnyIteratorListMappingOptional)
deprecated)CallbackManagerForLLMRun)BaseLLM)
GenerationGenerationChunk	LLMResult)
ConfigDictZgpt2text-generation)text2text-generationr   summarizationtranslation   z0.0.37z1.0z)langchain_huggingface.HuggingFacePipeline)ZsinceZremovalZalternative_importc                   @  s   e Zd ZU dZdZded< eZded< dZded< dZ	ded	< e
Zd
ed< eddZeddddde
fdddddddd
dd d
ddZeddddZeddddZd&ddddddd d!Zd'ddddd"d#d$d%ZdS )(HuggingFacePipelinea\  HuggingFace Pipeline API.

    To use, you should have the ``transformers`` python package installed.

    Only supports `text-generation`, `text2text-generation`, `summarization` and
    `translation`  for now.

    Example using from_model_id:
        .. code-block:: python

            from langchain_community.llms import HuggingFacePipeline
            hf = HuggingFacePipeline.from_model_id(
                model_id="gpt2",
                task="text-generation",
                pipeline_kwargs={"max_new_tokens": 10},
            )
    Example passing pipeline in directly:
        .. code-block:: python

            from langchain_community.llms import HuggingFacePipeline
            from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

            model_id = "gpt2"
            tokenizer = AutoTokenizer.from_pretrained(model_id)
            model = AutoModelForCausalLM.from_pretrained(model_id)
            pipe = pipeline(
                "text-generation", model=model, tokenizer=tokenizer, max_new_tokens=10
            )
            hf = HuggingFacePipeline(pipeline=pipe)
    Nr   pipelinestrmodel_idzOptional[dict]model_kwargspipeline_kwargsint
batch_sizeZforbid)extradefaultzOptional[int]zOptional[str])
r   taskbackenddevice
device_mapr   r   r   kwargsreturnc	              
   K  s  z$ddl m}
m}m} ddl m} W n ty>   tdY n0 |pFi }|j|fi |}zF|dkr|dkrzddlm} W n ty   tdY n0 z|j|fi |}W q t	y   |j|fd	d
i|}Y q0 n|
j|fi |}n|dv r|dkrvzddlm
} W n ty0   tdY n0 z|j|fi |}W n* t	yr   |j|fd	d
i|}Y n0 n|j|fi |}ntd| dt dW n8 ty } ztd| d|W Y d}~n
d}~0 0 |jdu rP|jjdur|jj|_nN|jjdur,t|jjtr,|jj|_n$|jdurB|j|_n|ddi t|ddslt|ddr|dur|dkrtd| d d}|dur:tjddur:|dkr:ddl}|j }|dk s||krtd| d| d|dur|dk rd}|dur:|dk r:|dkr:td | |durb|durb|dkrbtd! d"|v r~d#d$ | D }|pi }|f |||||||d%|}|jtvrtd|j dt d| f |||||d&|	S )'z5Construct the pipeline object from model_id and task.r   )AutoModelForCausalLMAutoModelForSeq2SeqLMAutoTokenizer)r   z`Could not import transformers python package. Please install it with `pip install transformers`.r   Zopenvino)OVModelForCausalLMzlCould not import optimum-intel python package. Please install it with: pip install 'optimum[openvino,nncf]' ZexportT)r   r   r   )OVModelForSeq2SeqLMGot invalid task , currently only  are supportedzCould not load the z# model due to missing dependencies.N	pad_tokenz[PAD]Zis_loaded_in_4bitFZis_loaded_in_8bitr   z+Setting the `device` argument to None from z to avoid the error caused by attempting to move the model that was already loaded on the GPU using the Accelerate module to the same or another device.torchr   zGot device==z', device is required to be within [-1, )zDevice has %d GPUs available. Provide device={deviceId} to `from_model_id` to use availableGPUs for execution. deviceId is -1 (default) for CPU and can be a positive integer associated with CUDA device id.z6Please set device for OpenVINO through: `model_kwargs`trust_remote_codec                 S  s   i | ]\}}|d kr||qS )r0    ).0kvr1   r1   {/var/www/html/cobodadashboardai.evdpl.com/venv/lib/python3.9/site-packages/langchain_community/llms/huggingface_pipeline.py
<dictcomp>   s   z5HuggingFacePipeline.from_model_id.<locals>.<dictcomp>)r   model	tokenizerr!   r"   r   r   )r   r   r   r   r   )transformersr%   r&   r'   r   ImportErrorZfrom_pretrainedZoptimum.intel.openvinor(   	Exceptionr)   
ValueErrorVALID_TASKSr-   configZpad_token_idZeos_token_id
isinstancer   Zadd_special_tokensgetattrloggerwarning	importlibutil	find_specr.   cudaZdevice_countitemsr   )clsr   r   r    r!   r"   r   r   r   r#   r%   r&   r'   Zhf_pipelineZ_model_kwargsr8   r(   r7   r)   er.   Zcuda_device_countZ_pipeline_kwargsr   r1   r1   r5   from_model_idM   s.   

















z!HuggingFacePipeline.from_model_idzMapping[str, Any])r$   c                 C  s   | j | j| jdS )zGet the identifying parameters.r   r   r   rK   selfr1   r1   r5   _identifying_params   s    z'HuggingFacePipeline._identifying_paramsc                 C  s   dS )NZhuggingface_pipeliner1   rL   r1   r1   r5   	_llm_type  s    zHuggingFacePipeline._llm_typez	List[str]zOptional[List[str]]z"Optional[CallbackManagerForLLMRun]r   )promptsstoprun_managerr#   r$   c                 K  s6  g }| j r| j ni }|d|}|dd}tdt|| jD ]}	||	|	| j  }
| j|
fi |}t|D ]\}}t|tr|d }| jj	dkr|d }n\| jj	dkr|d }nF| jj	dkr|d	 }n0| jj	d
v r|d }nt
d| jj	 dt d|r|t|
| d  }|| qnq>tdd |D dS )Nr   skip_promptFr   r   Zgenerated_textr   r   Zsummary_textr   Ztranslation_textr*   r+   r,   c                 S  s   g | ]}t |d gqS )text)r   )r2   rU   r1   r1   r5   
<listcomp>;      z1HuggingFacePipeline._generate.<locals>.<listcomp>)Zgenerations)r   getrangelenr   r   	enumerater?   listr   r<   r=   appendr   )rM   rP   rQ   rR   r#   Ztext_generationsZdefault_pipeline_kwargsr   rS   iZbatch_prompts	responsesjresponserU   r1   r1   r5   	_generate
  sB    




zHuggingFacePipeline._generatezIterator[GenerationChunk])promptrQ   rR   r#   r$   c                 +  s  ddl m} dd lddlm}m}m} |di }	|dd}
|d urV| jj	
|}|p\g  G  fddd|}|| g}| jj	|d	d
}|| jj	d|
dd}t|f||d|	}|| jjj|d}|  |D ](}t|d}|r|j|j|d |V  qd S )Nr   )Thread)StoppingCriteriaStoppingCriteriaListTextIteratorStreamerr   rS   Tc                      s$   e Zd Zddddd fddZdS )	z1HuggingFacePipeline._stream.<locals>.StopOnTokensztorch.LongTensorztorch.FloatTensorr   bool)	input_idsscoresr#   r$   c                   s$    D ]}|d d |kr dS qdS )Nr   r   TFr1   )rM   ri   rj   r#   Zstop_id)stopping_ids_listr1   r5   __call__V  s    z:HuggingFacePipeline._stream.<locals>.StopOnTokens.__call__N)__name__
__module____qualname__rl   r1   rk   r.   r1   r5   StopOnTokensU  s   rq   pt)Zreturn_tensorsg      N@)timeoutrS   Zskip_special_tokens)streamerstopping_criteria)targetr#   rT   )chunk)	threadingrd   r.   r9   re   rf   rg   rX   r   r8   Zconvert_tokens_to_idsdictr7   generatestartr   Zon_llm_new_tokenrU   )rM   rc   rQ   rR   r#   rd   re   rf   rg   r   rS   rq   ru   inputsrt   Zgeneration_kwargst1charrw   r1   rp   r5   _stream>  s@    
zHuggingFacePipeline._stream)NN)NN)rm   rn   ro   __doc__r   __annotations__DEFAULT_MODEL_IDr   r   r   DEFAULT_BATCH_SIZEr   r   Zmodel_configclassmethodrJ   propertyrN   rO   rb   r   r1   r1   r1   r5   r      s8   
$ 0  7  r   )
__future__r   importlib.utilrC   loggingtypingr   r   r   r   r   Zlangchain_core._api.deprecationr   Zlangchain_core.callbacksr	   Z#langchain_core.language_models.llmsr
   Zlangchain_core.outputsr   r   r   Zpydanticr   r   ZDEFAULT_TASKr=   r   	getLoggerrm   rA   r   r1   r1   r1   r5   <module>   s&   
