
    +#h7                        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
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mZ dZd	Zd
ZdZ ej6                  e      Z G d de      Zy)    )annotationsN)AnyDictIteratorListMappingOptional)CallbackManagerForLLMRun)BaseLLM)
GenerationGenerationChunk	LLMResult)
ConfigDictmodel_validatorgpt2text-generation)text2text-generationr   summarizationtranslation   c                  F   e Zd ZU dZdZded<   dZded<   	 dZded<   	 dZded	<   	 e	Z
d
ed<   	  ed      Z ed      edd              Zeddddde	f	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 dd       Zedd       Zedd       Z	 	 d	 	 	 	 	 	 	 	 	 ddZ	 	 d	 	 	 	 	 	 	 	 	 ddZy)HuggingFacePipelineaV  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_huggingface 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_huggingface 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Optional[str]model_idOptional[dict]model_kwargspipeline_kwargsint
batch_sizeforbid)extrabefore)modec                n    d|vr0d|v r#|d   r|d   j                   j                  |d<   |S t        |d<   |S )z8Ensure model_id is set either by pipeline or user input.r   r   )modelname_or_pathDEFAULT_MODEL_ID)clsvaluess     l/var/www/html/sandstorm/venv/lib/python3.12/site-packages/langchain_huggingface/llms/huggingface_pipeline.pypre_init_validatorz&HuggingFacePipeline.pre_init_validatorI   sN     V#V#z(:%+J%7%=%=%J%Jz"  &6z"    defaultc	           
     Z   	 ddl m}
m}m} ddl m} |r|j                         ni }|!|t        d      d|v rt        d      ||d<    |j                  |fi |}	 |d	k(  r5|d
k(  r	 ddl	m
} 	  |j                  |fi |}nc |
j                  |fi |}nO|dv r5|d
k(  r	 ddl	m} 	  |j                  |fi |}n* |j                  |fi |}nt        d| dt         d      |j                  |j                  j                   |_        t%        |dd      st%        |dd      r"| |dk(  rt&        j)                  d| d       d}|t*        j,                  j/                  d      j|dk(  reddl}|j2                  j5                         }|dk  s||k\  rt        d| d| d       ||dk  rd}| |dk  r|dkD  rt&        j)                  d!|       |||d
k(  rt&        j)                  d"       d#|v r)|j7                         D ci c]  \  }}|d#k7  s|| }}}|xs i } |d&||||||d$|}|j8                  t        vr t        d|j8                   dt         d       | d&|||||d%|	S # t
        $ r t        d      w xY w# t
        $ r t        d      w xY w# t        $ r  |j                  |fddi|}Y w xY w# t
        $ r t        d      w xY w# t        $ r  |j                  |fddi|}Y w xY w# t
        $ r}t        d| d      |d}~ww xY wc c}}w )'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`.NzBoth `device` and `device_map` are specified. `device` will override `device_map`. You will most likely encounter unexpected behavior.Please remove `device` and keep `device_map`.
device_mapz4`device_map` is already specified in `model_kwargs`.r   openvino)OVModelForCausalLMzlCould not import optimum-intel python package. Please install it with: pip install 'optimum[openvino,nncf]' exportT)r   r   r   )OVModelForSeq2SeqLMGot invalid task , currently only  are supportedzCould not load the z# model due to missing dependencies.is_loaded_in_4bitF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.torchz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_code)taskr&   	tokenizerdevicer    r   )r   r   r   r   r     )transformersr0   r1   r2   r   ImportError
ValueErrorcopyfrom_pretrainedoptimum.intel.openvinor5   	Exceptionr7   VALID_TASKS	pad_tokenconfigeos_token_idpad_token_idgetattrloggerwarning	importlibutil	find_specr=   cudadevice_countitemsrA   )r)   r   rA   backendrC   r3   r   r   r    kwargsr0   r1   r2   hf_pipeline_model_kwargsrB   r5   r&   r7   er=   cuda_device_countkv_pipeline_kwargsr   s                             r+   from_model_idz!HuggingFacePipeline.from_model_idT   s   	 
 = 0<))+!! $  }, !WXX*4M,'1M11(LmL	?	((j(

 B 2 B B$!(5! A0@@ $1E QQj(N
 C 3 C C$!(5! B1AA $1E !'v .&&1].B  &%*\\%>%>I" 2E:5"5u="9$NN=fX F" " F ((1=9$ %

 7 7 9{v):: "6( +<<M;NaQ  %&1*!fqj5F5JP & *"8W
=RNNST-/!.!4!4!6A!?R:R1M  +0b 
!&
 
 ==+#HMM? 3""-n>   
&,!
 
 	
o  	E 	> ' (D  %  B 2 B B$!-1!5B! ' (D  %  C 3 C C$!-1!5B!  	%dV+NO	ds   I4 
L (J /J$ L !K	 (K! ;+L 	L'L'4J	J!!L $KL KL 	KL !L?L LL 	L$LL$c                J    | j                   | j                  | j                  dS )zGet the identifying parameters.r   r   r   re   selfs    r+   _identifying_paramsz'HuggingFacePipeline._identifying_params
  s'      --#33
 	
r-   c                     y)Nhuggingface_pipelinerD   rf   s    r+   	_llm_typezHuggingFacePipeline._llm_type  s    %r-   c           	        g }|j                  di       }|j                  dd      }t        dt        |      | j                        D ]  }|||| j                  z    }	 | j                  |	fi |}
t        |
      D ]  \  }}t        |t              r|d   }| j                  j                  dk(  r|d   }n| j                  j                  dk(  r|d   }ng| j                  j                  dk(  r|d	   }nH| j                  j                  d
v r|d   }n*t        d| j                  j                   dt         d      |r|t        |	|         d  }|j                  |         t        |D cg c]  }t        |      g c}      S c c}w )Nr   skip_promptFr   r   generated_textr   r   summary_textr   translation_textr8   r9   r:   text)generations)getrangelenr    r   	enumerate
isinstancelistrA   rG   rL   appendr   r   )rg   promptsstoprun_managerr[   text_generationsr   rm   ibatch_prompts	responsesjresponserr   s                 r+   	_generatezHuggingFacePipeline._generate  s    ') **%6;jj6q#g,8 	.A#ADOO(;<M &!I  )3 .8h-'{H==%%)::#$45D]]''+AA#$45D]]''?:#N3D]]''=8#$67D$+DMM,>,>+? @**5nF  M!$4 5 78D ''-+.	.B =MNT*$/0N
 	
Ns   -F	c              +     K   ddl m} dd l}ddlm}m}m}	 |j                  di       }
|j                  dd      }|%| j                  j                  j                  |      }|xs g  G fdd|      } | |       g      } |	| j                  j                  d	|d
      }t        d|||d|
} || j                  |      }|j                          |D ]1  }t        |      }|r|j                  |j                  |       | 3 y w)Nr   )Thread)StoppingCriteriaStoppingCriteriaListTextIteratorStreamerr   rm   Tc                  *    e Zd Z	 	 	 	 	 	 	 	 d fdZy)1HuggingFacePipeline._stream.<locals>.StopOnTokensc                .    D ]  }|d   d   |k(  s y y)Nr   r>   TFrD   )rg   	input_idsscoresr[   stop_idstopping_ids_lists        r+   __call__z:HuggingFacePipeline._stream.<locals>.StopOnTokens.__call__`  s-      1 $G |B'72#$ r-   N)r   ztorch.LongTensorr   ztorch.FloatTensorr[   r   returnbool)__name__
__module____qualname__r   )r   s   r+   StopOnTokensr   _  s.    	+	 *	 		
 	r-   r   g      N@)timeoutrm   skip_special_tokens)text_inputsstreamerstopping_criteria)targetr[   rq   )chunkrD   )	threadingr   r=   rE   r   r   r   rt   r   rB   convert_tokens_to_idsdictstartr   on_llm_new_tokenrr   )rg   promptr|   r}   r[   r   r=   r   r   r   r   rm   r   r   r   generation_kwargst1charr   r   s                      @r+   _streamzHuggingFacePipeline._streamH  s     	%	
 	
 !**%6;jj5==**@@FD JB
	+ 
	 1,.1AB'MM### $	
 ! 
/
 	
 4==1BC

 	D#.E,,UZZu,EK	s   DD)r*   Dict[str, Any]r   r   )r   strrA   r   rZ   r   rC   zOptional[int]r3   r   r   r   r   r   r    r   r[   r   r   r   )r   zMapping[str, Any])r   r   )NN)
r{   z	List[str]r|   Optional[List[str]]r}   "Optional[CallbackManagerForLLMRun]r[   r   r   r   )
r   r   r|   r   r}   r   r[   r   r   zIterator[GenerationChunk])r   r   r   __doc__r   __annotations__r   r   r   DEFAULT_BATCH_SIZEr    r   model_configr   classmethodr,   rc   propertyrh   rk   r   r   rD   r-   r+   r   r      s   > Hc"Hm"B $(L.'0&*O^*3(J(HL (#  $ 
 ! $$('+*.,s
s
 s
 	s

 s
 "s
 %s
 (s
 s
 s
 
s
 s
j 
 
 & & %):>	/
/
 "/
 8	/

 /
 
/
h %):>	99 "9 8	9
 9 
#9r-   r   )
__future__r   importlib.utilrT   loggingtypingr   r   r   r   r   r	   langchain_core.callbacksr
   #langchain_core.language_models.llmsr   langchain_core.outputsr   r   r   pydanticr   r   r(   DEFAULT_TASKrL   r   	getLoggerr   rR   r   rD   r-   r+   <module>r      s]    "   ? ? = 7 I I 0    			8	$h' hr-   