o ZhtQ@sddlZddlZddlZddlmZddlmZmZm Z m Z ddl m Z m Z mZmZmZe r;ddlZddlmZmZerJddlZddlmZmZGd d d e Zeed d Gd ddeZdS)N)GenerationConfig)add_end_docstringsis_tf_availableis_torch_availablerequires_backends)ArgumentHandlerDatasetPipelinePipelineExceptionbuild_pipeline_init_args),MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES0MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES)/TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES3TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMESc@seZdZdZdddZdS)%TableQuestionAnsweringArgumentHandlerzB Handles arguments for the TableQuestionAnsweringPipeline NcKs`t|dddl}|durtd|durt|tr,|ddur,|ddur,|g}nbt|trmt|dkrmtdd|DsLtdd d|D|dddura|dddura|}n-td |d d t durvt|t s|t|t j r~|Std t |d ||dg}|D]}t|d|js|ddurtd||d|d<q|S)Npandasrz(Keyword argument `table` cannot be None.querytablecss|]}t|tVqdSN) isinstancedict.0dr^/var/www/auris/lib/python3.10/site-packages/transformers/pipelines/table_question_answering.py 6szATableQuestionAnsweringArgumentHandler.__call__..z:Keyword argument `table` should be a list of dict, but is css|]}t|VqdSr)typerrrrr8szIf keyword argument `table` is a list of dictionaries, each dictionary should have a `table` and `query` key, but only dictionary has keys z `table` and `query` keys.zZInvalid input. Keyword argument `table` should be either of type `dict` or `list`, but is ))rrzTable cannot be None.)rr ValueErrorrrgetlistlenallkeysr types GeneratorTyperZ DataFrame)selfrrkwargspdZtqa_pipeline_inputsZtqa_pipeline_inputrrr__call__&sD &$   z.TableQuestionAnsweringArgumentHandler.__call__)NN)__name__ __module__ __qualname____doc__r,rrrrr!srT)Z has_tokenizercszeZdZdZdZdZeddZeffdd Z dd Z d d Z fd d Z dddZ dddZdddZddZZS)TableQuestionAnsweringPipelinea Table Question Answering pipeline using a `ModelForTableQuestionAnswering`. This pipeline is only available in PyTorch. Unless the model you're using explicitly sets these generation parameters in its configuration files (`generation_config.json`), the following default values will be used: - max_new_tokens: 256 Example: ```python >>> from transformers import pipeline >>> oracle = pipeline(model="google/tapas-base-finetuned-wtq") >>> table = { ... "Repository": ["Transformers", "Datasets", "Tokenizers"], ... "Stars": ["36542", "4512", "3934"], ... "Contributors": ["651", "77", "34"], ... "Programming language": ["Python", "Python", "Rust, Python and NodeJS"], ... } >>> oracle(query="How many stars does the transformers repository have?", table=table) {'answer': 'AVERAGE > 36542', 'coordinates': [(0, 1)], 'cells': ['36542'], 'aggregator': 'AVERAGE'} ``` Learn more about the basics of using a pipeline in the [pipeline tutorial](../pipeline_tutorial) This tabular question answering pipeline can currently be loaded from [`pipeline`] using the following task identifier: `"table-question-answering"`. The models that this pipeline can use are models that have been fine-tuned on a tabular question answering task. See the up-to-date list of available models on [huggingface.co/models](https://huggingface.co/models?filter=table-question-answering). z table,queryT)Zmax_new_tokenscstj|i|||_|jdkrt}|tn t}|t | |t t |j jddo)r)inputsrrrbatch_inferencesz.TableQuestionAnsweringPipeline.batch_inferencec sh|jdkr`g}g}d}|djd}|d|j}|d|j}|d|j}d} t|D]} |dur| dddf} t| } || } t| jdD]@} | dddf | }| dddf | d}| ddd f | d}|dkr|dkr|dkrt |||f| | <qTt |  t j|j| dddf<|| }|| }|| } |j|d|d| dd }|j}|jr||j||t jj|d }|j| t j|jj}ttt| D]D\} }| dddf | }| dddf | d}| ddd f | d}|dkr<|dkr<|dkr<||f|qfd d D}q1t t|d}|jsV|fS|t t|dfSg}g}d}|djd}|d}|d}|d}d} t|D]} |dur| dddf} tj| tj d} || } t| jdD]D} | dddf | }| dddf | d}| ddd f | d}|dkr|dkr|dkrt |||f| | <q| | dddf<|| }|| }|| } |jtj!|ddtj!|ddtj!| ddd }|j}|jr!||j||t"j#$t"%|t"jt"%|t"j}tt| } tt"| D]E\} }| dddf | }| dddf | d}| ddd f | d}|dkr|dkr|dkr||f|qJfdd D}qt"&t|d}|js|fS|t"&t|dfS)z Inference used for models that need to process sequences in a sequential fashion, like the SQA models which handle conversational query related to a table. ptN input_idsrattention_masktoken_type_idsrr)rIrJrK)logitsc$i|]}|t|dkqSg?nparraymeanrkeyZcoords_to_probsrr $zGTableQuestionAnsweringPipeline.sequential_inference..)Zdtype)ZaxiscrNrOrPrTrVrrrWrX)'r9shapetoZdevicerangerQZ zeros_likecpunumpytolistinttorchZ from_numpyrlongr>Z unsqueezerMr@appendZlogits_aggregation distributionsZ BernoulliZprobsZfloat32 collections defaultdictr# enumerateZsqueezecattupleZint32Z expand_dimsr3mathZsigmoidcastconcat)r)rFZ all_logitsZall_aggregationsZ prev_answersZ batch_sizerIrJrKZtoken_type_ids_exampleindexZprev_labels_exampleZ model_labelsiZ segment_idZcol_idZrow_idZinput_ids_exampleZattention_mask_exampleoutputsrMZdist_per_tokenZ probabilitiespcolrowZ logits_batchrrVrsequential_inferences &   "        "z3TableQuestionAnsweringPipeline.sequential_inferencecs<|j|i|}tj|fi|}t|dkr|dS|S)a Answers queries according to a table. The pipeline accepts several types of inputs which are detailed below: - `pipeline(table, query)` - `pipeline(table, [query])` - `pipeline(table=table, query=query)` - `pipeline(table=table, query=[query])` - `pipeline({"table": table, "query": query})` - `pipeline({"table": table, "query": [query]})` - `pipeline([{"table": table, "query": query}, {"table": table, "query": query}])` The `table` argument should be a dict or a DataFrame built from that dict, containing the whole table: Example: ```python data = { "actors": ["brad pitt", "leonardo di caprio", "george clooney"], "age": ["56", "45", "59"], "number of movies": ["87", "53", "69"], "date of birth": ["7 february 1967", "10 june 1996", "28 november 1967"], } ``` This dictionary can be passed in as such, or can be converted to a pandas DataFrame: Example: ```python import pandas as pd table = pd.DataFrame.from_dict(data) ``` Args: table (`pd.DataFrame` or `Dict`): Pandas DataFrame or dictionary that will be converted to a DataFrame containing all the table values. See above for an example of dictionary. query (`str` or `List[str]`): Query or list of queries that will be sent to the model alongside the table. sequential (`bool`, *optional*, defaults to `False`): Whether to do inference sequentially or as a batch. Batching is faster, but models like SQA require the inference to be done sequentially to extract relations within sequences, given their conversational nature. padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`): Activates and controls padding. Accepts the following values: - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single sequence if provided). - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum acceptable input length for the model if that argument is not provided. - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different lengths). truncation (`bool`, `str` or [`TapasTruncationStrategy`], *optional*, defaults to `False`): Activates and controls truncation. Accepts the following values: - `True` or `'drop_rows_to_fit'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum acceptable input length for the model if that argument is not provided. This will truncate row by row, removing rows from the table. - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths greater than the model maximum admissible input size). Return: A dictionary or a list of dictionaries containing results: Each result is a dictionary with the following keys: - **answer** (`str`) -- The answer of the query given the table. If there is an aggregator, the answer will be preceded by `AGGREGATOR >`. - **coordinates** (`List[Tuple[int, int]]`) -- Coordinates of the cells of the answers. - **cells** (`List[str]`) -- List of strings made up of the answer cell values. - **aggregator** (`str`) -- If the model has an aggregator, this returns the aggregator. rr)r8r6r,r$)r)rBr*Zpipeline_inputsresultsrDrrr,s K z'TableQuestionAnsweringPipeline.__call__NcKsi}|dur ||d<|dur||d<i}|dur||d<t|dddur)|j|d<t|dddur;|j|d<|j|d<||ifS)Npadding truncation sequentialassistant_modelassistant_tokenizer tokenizer)r=rwryrx)r)rvrtrur*Zpreprocess_paramsZforward_paramsrrr_sanitize_parametersjs    z3TableQuestionAnsweringPipeline._sanitize_parameterscCsv|dur|jdkr d}nd}|d|d}}|jrtd|dus&|dkr*td|j|||j||d }||d<|S) Nr5Zdrop_rows_to_fitZdo_not_truncaterrztable is emptyzquery is empty)Zreturn_tensorsrurt)remptyr!ryr9)r)Zpipeline_inputrvrtrurrrFrrr preprocess}s z)TableQuestionAnsweringPipeline.preprocessFcKst|d}|jdkr|r|jdi|}n|jdi|}nd|vr'|j|d<|jjdi||}|||d}|S)Nrr5generation_config) model_inputsrrnr)poprrrrGr~r>generate)r)rrvZgenerate_kwargsrrn model_outputsrrr_forwards    z'TableQuestionAnsweringPipeline._forwardcsr|d}|d|d}jdkrjrE|dd\}}j|||}|\}}fddt|Djjjfddt|D} n|d }j||}|d }ii} g} t|D]6\} } fd d | D} | d }| | d }|d | | fdd | Dd}|r||d<| |q^t |d krt dn dd jj |ddD} t | dkr| S| d S)Nrrrnr5rcs i|] \}}|jjj|qSr)r>r?r4rrmpred)r)rrrWs z>TableQuestionAnsweringPipeline.postprocess..cs&i|]\}}|kr||dqS)z > rr) aggregatorsno_agg_label_indexrrrWs rcg|]}j|qSrZiatrZ coordinaterrr z>TableQuestionAnsweringPipeline.postprocess..r{z, crrrrrrrrr)answer coordinatescells aggregatorz Empty answercSsg|]}d|iqS)rr)rrrrrrsT)Zskip_special_tokensr)rr@ryZconvert_logits_to_predictionsrfr>r?Zno_aggregation_label_indexr"joinrbr$r Z batch_decode)r)rrFrnrMZ logits_aggZ predictionsZanswer_coordinates_batchZagg_predictionsZaggregators_prefixZanswersrlrrrZaggregator_prefixrr)rrr)rr postprocesssF        z*TableQuestionAnsweringPipeline.postprocess)NNN)NTN)F)r-r.r/r0Zdefault_input_namesZ_pipeline_calls_generaterZ_default_generation_configrr7rGrrr,rzr}rr __classcell__rrrDrr1Vs "  R  r1)rdr'r]rQZ generationrutilsrrrrbaser r r r r r`Zmodels.auto.modeling_autorrZ tensorflowr3Zmodels.auto.modeling_tf_autorrrr1rrrrs  5