o Zhw|@sddlZddlZddlZddlZddlZddlZddlmZmZddl m Z m Z ddl m Z ddlmZddlZddlmZmZddlmZed krPddlZeredd lmZdd lmZdd lmZerzddlZdd lm Z m!Z!m"Z"m#Z#m$Z$e%ej&ej'Z(e%ej&ej)ej'dZ*ddiddiddiddiddiddidZ+dZ,dd-e+.dZ/gdZ0GdddZ1e GdddZ2d efd!d"Z3Gd#d$d$eZ4dS)%N)ArgumentParser Namespace) dataclassfield)Thread)Optional)is_rich_availableis_torch_available)BaseTransformersCLICommandWindows)Console)Live)Markdown)AutoModelForCausalLM AutoTokenizerBitsAndBytesConfigGenerationConfigTextIteratorStreamerz .!\"#$%&'()*+,\-/:<=>?@[]^_`{|}~textz5There is a Llama in my lawn, how can I get rid of it?zyWrite a Python function that integrates any Python function f(x) numerically over an arbitrary interval [x_start, x_end].z4How many helicopters can a human eat in one sitting?z4Count to 10 but skip every number ending with an 'e'zWhy aren't birds real?z2Why is it important to eat socks after meditating?)llamacode helicopternumbersZbirdssocksaa **TRANSFORMERS CHAT INTERFACE** Chat interface to try out a model. Besides chatting with the model, here are some basic commands: - **!help**: shows all available commands - **!status**: shows the current status of the model and generation settings - **!clear**: clears the current conversation and starts a new one - **!exit**: closes the interface am **TRANSFORMERS CHAT INTERFACE HELP** Full command list: - **!help**: shows this help message - **!clear**: clears the current conversation and starts a new one - **!status**: shows the current status of the model and generation settings - **!example {NAME}**: loads example named `{NAME}` from the config and uses it as the user input. Available example names: `z`, `a%` - **!set {ARG_1}={VALUE_1} {ARG_2}={VALUE_2}** ...: changes the system prompt or generation settings (multiple settings are separated by a space). Accepts the same flags and format as the `generate_flags` CLI argument. If you're a new user, check this basic flag guide: https://huggingface.co/docs/transformers/llm_tutorial#common-options - **!save {SAVE_NAME} (optional)**: saves the current chat and settings to file by default to `./chat_history/{MODEL_NAME}/chat_{DATETIME}.yaml` or `{SAVE_NAME}` if provided - **!exit**: closes the interface ) )max_new_tokensr) do_sampleTr) num_beamsr r) temperature?r)top_k2r!)top_pr r#)repetition_penaltyr r$) eos_tokensN eos_token_id) eos_token_idsNr&c@seZdZddeedeefddZdedefdd Zdefd d Zd d Z defddZ dedefddZ dde fddZ dededefddZdS) RichInterfaceN model_name user_namecCs:t|_|dur d|_n||_|durd|_dS||_dS)N assistantuser)r _consoler)r*)selfr)r*r/I/var/www/auris/lib/python3.10/site-packages/transformers/commands/chat.py__init__xs  zRichInterface.__init__ output_streamreturnc Csd}|jd|jdt|jddJ}t|D]=\}}|r#|dkr$q||7}g}|D]}|||dr@|dq.|d q.td | d d }| |qWd n1sawY|j|S) zRStream output from a role, and return the generated text after it's done steaming.z [bold blue]:)consolerefresh_per_secondrz``` z z github-dark)Z code_themeN) r-printr)r enumerate splitlinesappend startswithrjoinstripupdate) r.r2rliveiZoutputslineslinemarkdownr/r/r0 stream_outputs&        zRichInterface.stream_outputcCs$|jd|jd}|j|S)z!Gets user input from the console. [bold red]: )r-inputr*r9)r.rHr/r/r0rHs zRichInterface.inputcCs|jdS)zClears the console.N)r-clear)r.r/r/r0rIszRichInterface.clearrcCs(|jd|jd||jdS)z%Prints a user message to the console.rGz>:[/ bold red] N)r-r9r*)r.rr/r/r0print_user_messagesz RichInterface.print_user_messagecolorcCs&|jd|d||jdS)z,Prints text in a given color to the console.z[bold ]Nr-r9)r.rrKr/r/r0 print_colorzRichInterface.print_colorFminimalcCs&|jt|rtnt|jdS)z'Prints the help message to the console.N)r-r9rHELP_STRING_MINIMAL HELP_STRING)r.rPr/r/r0 print_helprOzRichInterface.print_helpgeneration_config model_kwargscCsJ|jd|d|r|jd||jd||jdS)zFPrints the status of the model and generation settings to the console.z[bold blue]Model: r8z[bold blue]Model kwargs: z [bold blue]NrM)r.r)rTrUr/r/r0 print_statuss zRichInterface.print_status)NN)F)__name__ __module__ __qualname__rstrr1rrFrHrIrJrNboolrSrdictrVr/r/r/r0r(ws 'r(c@seZdZUdZedddidZeeed<edddidZ eeed<eddd idZ eeed <ed dd idZ eed <edddidZ eeed<edddidZ eeed<edddidZeed<edddidZeed<edddidZeed<edddidZeed<edddidZeed <eddd!idZeed"<eddd#idZeed$<eddd%idZeeed&<eddd'idZeeed(<ed)dd*idZeed+<ed,dd-idZeed.<ed/d0gd1d2dZeeed3<ed4dd5idZeed6<eddd7idZeeed8<ed4dd9idZeed:<ed4dd;idZ eed<<ed=d>d?d=gd2dZ!eed@<ed4ddAidZ"eedB<dS)C ChatArgumentsz Arguments for the chat CLI. See the metadata arg for each argument's description -- the medatata will be printed with `transformers chat --help` Nhelpz_Name of the pre-trained model. The positional argument will take precedence if both are passed.)defaultmetadatamodel_name_or_pathzKUsername to display in chat interface. Defaults to the current user's name.r,zSystem prompt. system_promptz./chat_history/zFolder to save chat history. save_folderz"Path to a yaml file with examples. examples_pathzPath to a local generation config file or to a HuggingFace repo containing a `generation_config.json` file. Other generation settings passed as CLI arguments will be applied on top of this generation config.rTrz%Maximum number of tokens to generate.rTz,Whether to sample outputs during generation.rr z Number of beams for beam search.rr z%Temperature parameter for generation.rr"zValue of k for top-k sampling.r!z Value of p for nucleus sampling.r#zRepetition penalty.r$zNEOS tokens to stop the generation. If multiple they should be comma separated.r%zQEOS token IDs to stop the generation. If multiple they should be comma separated.r'mainzLSpecific model version to use (can be a branch name, tag name or commit id).model_revisioncpuzDevice to use for inference.deviceautozOverride the default `torch.dtype` and load the model under this dtype. If `'auto'` is passed, the dtype will be automatically derived from the model's weights.)riZbfloat16Zfloat16Zfloat32)r^choices torch_dtypeFz2Whether to trust remote code when loading a model.trust_remote_codezWhich attention implementation to use; you can run --attn_implementation=flash_attention_2, in which case you must install this manually by running `pip install flash-attn --no-build-isolation`.attn_implementationzIWhether to use 8 bit precision for the base model - works only with LoRA. load_in_8bitzIWhether to use 4 bit precision for the base model - works only with LoRA. load_in_4bitZnf4zQuantization type.Zfp4bnb_4bit_quant_typez#Whether to use nested quantization.use_bnb_nested_quant)#rWrXrY__doc__rrarrZ__annotations__r,rbrcrdrTrintrr[rrfloatr!r#r$r%r'rfrhrkrlrmrnrorprqr/r/r/r0r]s  r]argscCst|S)z; Factory function used to chat with a local model. ) ChatCommand)rvr/r/r0chat_command_factory%srxc@sxeZdZedefddZddZdedefdd Zede fd d Z ed-ded e e de fddZ ed-de e de efddZde e defddZdededeeeffddZedede e de e deee effddZedede dfdd Zdedeeeffd!d"Zd#e ded$ed%ee ee e ffd&ed'ed(e edee eeeffd)d*Zd+d,Zd S).rwparsercCsTtf}|jd|d}|d}|jdtddd|jdtdd d d |jtd dS) z Register this command to argparse so it's available for the transformer-cli Args: parser: Root parser to register command-specific arguments chat)dataclass_typeszPositional argumentsmodel_name_or_path_positionalNzName of the pre-trained model.)typer_r^generate_flagsaFlags to pass to `generate`, using a space as a separator between flags. Accepts booleans, numbers, and lists of integers, more advanced parameterization should be set through --generation-config. Example: `transformers chat max_new_tokens=100 do_sample=False eos_token_id=[1,2]`. If you're a new user, check this basic flag guide: https://huggingface.co/docs/transformers/llm_tutorial#common-options*)r}r_r^nargs)func)r] add_parseradd_argument_group add_argumentrZ set_defaultsrx)ryr{Z chat_parsergroupr/r/r0register_subcommand-s  zChatCommand.register_subcommandcCs||}||_dSN)_handle_deprecated_argsrv)r.rvr/r/r0r1Js  zChatCommand.__init__rvr3c Csd}|jp|j|_|jdurtd|jdurd}tdttD] \}}}t||}||krAd}td|d|d|d tq!|rKtd t |S) z Handles deprecated arguments and their deprecation cycle. To be removed after we fully migrated to the new args. FNzOne of the following must be provided: - The positional argument containing the model repo, e.g. `transformers chat ` - the optional --model_name_or_path argument, containing the model repo (deprecated)TzThe --model_name_or_path argument is deprecated will be removed in v4.54.0. Use the positional argument instead, e.g. `transformers chat `.zThe --a. argument is deprecated will be removed in v4.54.0. There are two alternative solutions to specify this generation option: 1. Pass `--generation-config ` to specify a generation config. 2. Pass `generate` flags through positional arguments, e.g. `transformers chat =`z (Press enter to continue)) r|ra ValueErrorwarningswarn FutureWarning_DEPRECATION_MAPgetattrr9rH)r.rvZ has_warningsZdeprecated_arg default_valueZnew_argvaluer/r/r0rNs:    z#ChatCommand._handle_deprecated_argscCs$tdkr tSttjS)z)Returns the username of the current user.r )platformsystemosgetloginpwdgetpwuidgetuidpw_namer/r/r/r0 get_usernameys zChatCommand.get_usernameNfilenamecCsi}t||d<||d<|j}|dur(td}|jd|d}tj||}tjtj |ddt |d }t j ||d d Wdn1sKwYtj |S) z!Saves the chat history to a file.settingsZ chat_historyNz%Y-%m-%d_%H-%M-%Sz/chat_.jsonT)exist_okwr5)indent)varsrctimestrftimer|rpathr>makedirsdirnameopenjsondumpabspath)rzrvrZ output_dictfolderZtime_strfr/r/r0 save_chats    zChatCommand.save_chatrbcCs |durg}|Sd|dg}|S)zClears the chat history.NrZrolecontentr/)rbrzr/r/r0clear_chat_historys  zChatCommand.clear_chat_historyr~cst|dkriSdd|D}dd|D}dd|D}dtdtfdd fd d|D}d d d |D}d|d}|dd}|dd}|dd}|dd}|dd}|dd}zt|}W|Stjyt dw)zUParses the generate flags from the user input into a dictionary of `generate` kwargs.rcSs.i|]}d|ddd|ddqS)"rrr )split).0flagr/r/r0 s.z4ChatCommand.parse_generate_flags..cSs*i|]\}}||dvr|n|qS))truefalse)lowerrkvr/r/r0rscSs"i|] \}}||dkr dn|qS)Nonenullr/rr/r/r0rs"sr3cSs|dddS)N.r4r )replaceisdigit)rr/r/r0 is_numbersz3ChatCommand.parse_generate_flags..is_numbercs*i|]\}}||sd|dn|qS)rr/rrr/r0rs*z, cSsg|] \}}|d|qS)z: r/rr/r/r0 sz4ChatCommand.parse_generate_flags..{}z"null"rz"true"rz"false"rz"[[z]"rLr:zFailed to convert `generate_flags` into a valid JSON object. `generate_flags` = {generate_flags} Converted JSON string = {generate_flags_string}) lenitemsrZr[r>rrloadsJSONDecodeErrorr)r.r~Zgenerate_flags_as_dictZgenerate_flags_stringZprocessed_generate_flagsr/rr0parse_generate_flagss2         z ChatCommand.parse_generate_flags tokenizerc Cs|jdur/t}|||j|j\}}|j|j|j|j|j |j |j ||d }|j di|n d|jvrIt j|j}t j|j}t||}nt|j}||j} |j di| } || fS)zj Returns a GenerationConfig object holding the generation parameters for the CLI command. N) rrrrr!r#r$ pad_token_idr&rr/)rTrparse_eos_tokensr%r'rrrrr!r#r$r@rrrbasenamefrom_pretrainedrr~) r.rvrrTrr'Zdeprecated_kwargsrrZparsed_generate_flagsrUr/r/r0get_generation_parameterizations,    z+ChatCommand.get_generation_parameterizationr%r'cCs||jdur |j}n|j}g}|dur|||d|dur.|dd|dDt|dkr:||j||fS)z:Retrieves the pad token ID and all possible EOS token IDs.N,cSsg|]}t|qSr/)rt)rZtoken_idr/r/r0r sz0ChatCommand.parse_eos_tokens..r)rr&extendZconvert_tokens_to_idsrrr<)rr%r'rZall_eos_token_idsr/r/r0rs   zChatCommand.parse_eos_tokens model_argsrcCs@|jrtd|j|j|j|jd}|S|jrtdd}|Sd}|S)NT)roZbnb_4bit_compute_dtyperpZbnb_4bit_use_double_quantZbnb_4bit_quant_storage)rn)rorrkrprqrn)rquantization_configr/r/r0get_quantization_configs z#ChatCommand.get_quantization_configcCstj|j|j|jd}|jdvr|jntt|j}||}|j|j |d|d}t j|jfd|ji|}t|dddurC| |j }||fS)N)revisionrl)riNri)rrmrkZ device_maprrlZ hf_device_map) rrr|rfrlrkrtorchrrmrtorh)r.rvrrkrrUmodelr/r/r0load_model_and_tokenizer&s.  z$ChatCommand.load_model_and_tokenizer user_input interfaceexamplesrTrUrzcCs|dkr||j}|n|dkr|n|drIt|dkrI|}t|dkr4|d} nd} |||| } |jd| dd d n|d r|d d } | } | D]} d | vro|jd| ddd nq\| | } |j di| } |j di| ne|drt|dkr|d}||vr|g}| ||d| d||ddn2d|dt|d}|j|dd n|dkr|j|j||dn|jd|ddd ||||fS)z Handles all user commands except for `!exit`. May update the chat history (e.g. reset it) or the generation config (e.g. set a new flag). z!clearz!helpz!saver NzChat saved in !green)rrKz!setr5rz(Invalid flag format, missing `=` after `z;`. Please use the format `arg_1=value_1 arg_2=value_2 ...`.red!examplerr,rzExample z* not found in list of available examples: rz!status)r)rTrU'z/' is not a valid command. Showing help message.r/)rrbrIrSr=rrrrNr?rr@rJr<listkeysrVr|)r.rrvrrrTrUrzZ split_inputrZnew_generate_flagsrZparsed_new_generate_flagsZnew_model_kwargsZ example_nameZ example_errorr/r/r0handle_non_exit_user_commandsAs\          z)ChatCommand.handle_non_exit_user_commandsc Cststdtstd|j}|jdurt}nt|j }t|}Wdn1s.wY|j dur=| }n|j }| |\}}t |ddd}| ||\}} t|j|d} | ||j} | jdd zg| } | dr| dkrWdS|j| || ||| | d \} }} | d sWqmn| d | d |j| d dd|j} t| }| |||d| }t|j|d}|| |}|!| d|d Wn t"yYdSwqn)NzHYou need to install rich to use the chat interface. (`pip install rich`)zJYou need to install torch to use the chat interface. (`pip install torch`)T)Zskip_special_tokensZ skip_prompt)r)r*)rPrz!exit)rrvrrrTrUrzrr,rpt)Zreturn_tensorsZadd_generation_prompt)inputsattention_maskstreamerrT)targetkwargsr+)#r ImportErrorr rvrdDEFAULT_EXAMPLESryamlZ safe_loadr,rrrrr(r|rIrrbrSrHr=rr<Zapply_chat_templaterrhrZ ones_likergeneratestartrFr>KeyboardInterrupt)r.rvrrr,rrZgeneration_streamerrTrUrrzrrrZgeneration_kwargsthreadZ model_outputr/r/r0runsv            zChatCommand.runr)rWrXrY staticmethodrrr1r]rrZrrrrr\rrrtuplerrrtrrrrr(rrr/r/r/r0rw,sd+ 3  '  Rrw)5rrrstringrrargparserr dataclassesrr threadingrtypingrrZtransformers.utilsrr r4r rrZ rich.consoler Z rich.liverZ rich.markdownrrZ transformersrrrrrset ascii_letters whitespaceZALLOWED_KEY_CHARSdigitsZALLOWED_VALUE_CHARSrrQr>rrRrr(r]rxrwr/r/r/r0sZ          UX