o Zh"@spddlZddlmZmZmZmZddlmZddlm Z ddl m Z ddl m Z dZdZd ZGd d d e ZdS) N)AnyListMappingOptional)CallbackManagerForLLMRun)LLM) ConfigDict)enforce_stop_tokenszgoogle/flan-t5-largetext2text-generation)r text-generation summarizationc@seZdZUdZdZeed<eZe ed< dZ e e ed< dZ e e ed< eddZe    dde d e d e ed e e de e de e de ede ede ededefddZedee effddZede fddZ  dde de ee de edede f ddZdS) WeightOnlyQuantPipelineaWeight only quantized model. To use, you should have the `intel-extension-for-transformers` packabge and `transformers` package installed. intel-extension-for-transformers: https://github.com/intel/intel-extension-for-transformers Example using from_model_id: .. code-block:: python from langchain_community.llms import WeightOnlyQuantPipeline from intel_extension_for_transformers.transformers import ( WeightOnlyQuantConfig ) config = WeightOnlyQuantConfig hf = WeightOnlyQuantPipeline.from_model_id( model_id="google/flan-t5-large", task="text2text-generation" pipeline_kwargs={"max_new_tokens": 10}, quantization_config=config, ) Example passing pipeline in directly: .. code-block:: python from langchain_community.llms import WeightOnlyQuantPipeline from intel_extension_for_transformers.transformers import ( AutoModelForSeq2SeqLM ) from intel_extension_for_transformers.transformers import ( WeightOnlyQuantConfig ) from transformers import AutoTokenizer, pipeline model_id = "google/flan-t5-large" tokenizer = AutoTokenizer.from_pretrained(model_id) config = WeightOnlyQuantConfig model = AutoModelForSeq2SeqLM.from_pretrained( model_id, quantization_config=config, ) pipe = pipeline( "text-generation", model=model, tokenizer=tokenizer, max_new_tokens=10, ) hf = WeightOnlyQuantPipeline(pipeline=pipe) Npipelinemodel_id model_kwargspipeline_kwargsZallow)extraFtaskdevice device_map load_in_4bit load_in_8bitquantization_configkwargsreturnc  Ks|durt|tr|dkrtdtjddurtdzddlm} m} ddl m } dd l m }dd l m }Wn tyDtd wt|tr\|dkr\| sUtd d t|}n t|trg|dkrgd}|durq|durqd}|pti}|j|fi|}z5|dkr| j|f||| d|d|}n|dvr| j|f||| d|d|}n td|dtdWnty}z td|d|d}~wwd|vrdd|D}|pi}|d|||||d|}|jtvrtd|jdtd|d||||d| S)z5Construct the pipeline object from model_id and task.Nrz7`Device` and `device_map` cannot be set simultaneously!Ztorchz;Weight only quantization pipeline only support PyTorch now!r)AutoModelForCausalLMAutoModelForSeq2SeqLM)is_ipex_available) AutoTokenizer)rzCould not import transformers python package. Please install it with `pip install transformers` and `pip install intel-extension-for-transformers`.z)Don't find out Intel GPU on this machine!zxpu:cpur F)rrrZuse_llm_runtimer)r r Got invalid task , currently only  are supportedzCould not load the z# model due to missing dependencies.trust_remote_codecSsi|] \}}|dkr||qS)r$).0kvr%r%h/var/www/html/lang_env/lib/python3.10/site-packages/langchain_community/llms/weight_only_quantization.py sz9WeightOnlyQuantPipeline.from_model_id..)rmodel tokenizerrr)rrrrr%) isinstanceint ValueError importlibutil find_specZ-intel_extension_for_transformers.transformersrrZ,intel_extension_for_transformers.utils.utilsrZ transformersrr ImportErrorstrZfrom_pretrained VALID_TASKSitemsr)clsrrrrrrrrrrrrrrZ hf_pipelineZ _model_kwargsr,r+eZ_pipeline_kwargsrr%r%r) from_model_idOs        z%WeightOnlyQuantPipeline.from_model_idcCs|j|j|jdS)zGet the identifying parameters.rrrr:selfr%r%r)_identifying_paramssz+WeightOnlyQuantPipeline._identifying_paramscCsdS)zReturn type of llm.Zweight_only_quantizationr%r;r%r%r) _llm_typesz!WeightOnlyQuantPipeline._llm_typepromptstop run_managercKs||}|jjdkr|ddt|d}n'|jjdkr%|dd}n|jjdkr2|dd}n td|jjd td |rFt||}|S) abCall the HuggingFace model and return the output. Args: prompt: The prompt to use for generation. stop: A list of strings to stop generation when encountered. Returns: The generated text. Example: .. code-block:: python from langchain_community.llms import WeightOnlyQuantPipeline llm = WeightOnlyQuantPipeline.from_model_id( model_id="google/flan-t5-large", task="text2text-generation", ) llm.invoke("This is a prompt.") r rZgenerated_textNr r Z summary_textr!r"r#)rrlenr/r5r )r<r?r@rArresponsetextr%r%r)_calls      zWeightOnlyQuantPipeline._call)rNNNFFN)NN)__name__ __module__ __qualname____doc__rr__annotations__DEFAULT_MODEL_IDrr4rrdictrrZ model_config classmethodr.boolrr9propertyrr=r>rrrEr%r%r%r)r sx  1     j r )r0typingrrrrZ langchain_core.callbacks.managerrZ#langchain_core.language_models.llmsrZpydanticrZlangchain_community.llms.utilsr rKZ DEFAULT_TASKr5r r%r%r%r)s