model_inference_pudf.ipynb

{"cells":[{"cell_type":"markdown","source":["### Multi-node distributed inference of *Hugging Face* models with *Transformers Pipeline* and *Pandas UDF*\n\n##### Here we perform model inference using the NLP model from [*Hugging Face*](https://huggingface.co/), after we fine tune it. We use [*Transformers Pipeline*](https://huggingface.co/docs/transformers/main_classes/pipelines) to perform the inference task and [*Pandas user-defined functions*](https://docs.microsoft.com/en-us/azure/databricks/spark/latest/spark-sql/udf-python-pandas) on *Spark* to provide the distributed infrastructure to parallelize the process."],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"e44de03f-6c5a-49d3-92a7-12393ade555e"}}},{"cell_type":"markdown","source":["Import the necessary packages."],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"80f41b17-a142-456e-83e7-ce07fd6c1c4a"}}},{"cell_type":"code","source":["import os\nimport time\n\nimport pandas as pd\n\nfrom pyspark.sql.functions import pandas_udf\nfrom pyspark.sql.types import FloatType, IntegerType, StructType, StructField\n\nimport torch\nfrom torch.utils.data import DataLoader, SequentialSampler\n\nfrom transformers import pipeline\nfrom transformers import AutoConfig, AutoTokenizer, AutoModelForSequenceClassification"],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"d5f4510b-bf8b-4297-83e7-c3d6189686b4"}},"outputs":[{"output_type":"display_data","metadata":{"application/vnd.databricks.v1+output":{"datasetInfos":[],"data":{"text/plain":"","application/vnd.databricks.v1+bamboolib_hint":"{\"pd.DataFrames\": [], \"version\": \"0.0.1\"}"},"removedWidgets":[],"addedWidgets":{},"metadata":{},"type":"mimeBundle","arguments":{}}},"output_type":"display_data","data":{"text/plain":"","application/vnd.databricks.v1+bamboolib_hint":"{\"pd.DataFrames\": [], \"version\": \"0.0.1\"}"}}],"execution_count":0},{"cell_type":"markdown","source":["Define the number of workers in the cluster and the folder in the Databricks File System to load the prepared data from."],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"1de88aa2-e7e1-4d6d-84cd-19793352a332"}}},{"cell_type":"code","source":["num_workers = 8\n\ntest_data_folder = '/test_data'"],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"02539113-1b89-4131-9de5-20fa4515bbea"}},"outputs":[],"execution_count":0},{"cell_type":"markdown","source":["Read the testing data into a *Spark DataFrame*. We repartition the data by the number of workers in the cluster. We do this to make sure all workers in the cluster are used to process the data and to illustrate the parallel processing with *Pandas UDF* below, given that the data is relatively small."],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"180e3bdd-de76-4a89-8db5-f46977ff73a4"}}},{"cell_type":"code","source":["sdf_test = spark.read.parquet(os.path.join(test_data_folder, 'test_data')).repartition(num_workers)\ndisplay(sdf_test.select(sdf_test.columns).limit(10))"],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"79f4c7ba-5b4d-415f-9958-17f924aa77bc"}},"outputs":[{"output_type":"display_data","metadata":{"application/vnd.databricks.v1+output":{"overflow":false,"datasetInfos":[],"data":[["a charming , banter-filled comedy . . . one of those airy cinematic bon bons whose aims -- and by extension , accomplishments -- seem deceptively slight on the surface .",1,60129542207],["the attempt to build up a pressure cooker of horrified awe emerges from the simple fact that the movie has virtually nothing to show .",0,60129542175],["marinated in clichés and mawkish dialogue .",0,60129542187],["a dark comedy that goes for sick and demented humor simply to do so . the movie is without intent .",0,60129542230],["the world needs more filmmakers with passionate enthusiasms like martin scorsese . but it doesn't need gangs of new york .",0,60129542260],["as his circle of friends keeps getting smaller one of the characters in long time dead says 'i'm telling you , this is f * * * ed' . maybe he was reading the minds of the audience .",0,60129542214],["scherfig's light-hearted profile of emotional desperation is achingly honest and delightfully cheeky .",1,60129542263],["the plot is romantic comedy boilerplate from start to finish .",0,60129542188],["humor in i spy is so anemic .",0,60129542206],["the increasingly diverse french director has created a film that one can honestly describe as looking , sounding and simply feeling like no other film in recent history .",1,60129542242]],"plotOptions":{"displayType":"table","customPlotOptions":{},"pivotColumns":null,"pivotAggregation":null,"xColumns":null,"yColumns":null},"columnCustomDisplayInfos":{},"aggType":"","isJsonSchema":true,"removedWidgets":[],"aggSchema":[],"schema":[{"name":"text","type":"\"string\"","metadata":"{}"},{"name":"label","type":"\"long\"","metadata":"{}"},{"name":"id","type":"\"long\"","metadata":"{}"}],"aggError":"","aggData":[],"addedWidgets":{},"metadata":{},"dbfsResultPath":null,"type":"table","aggOverflow":false,"aggSeriesLimitReached":false,"arguments":{}}},"output_type":"display_data","data":{"text/html":["<style scoped>\n  .table-result-container {\n    max-height: 300px;\n    overflow: auto;\n  }\n  table, th, td {\n    border: 1px solid black;\n    border-collapse: collapse;\n  }\n  th, td {\n    padding: 5px;\n  }\n  th {\n    text-align: left;\n  }\n</style><div class='table-result-container'><table class='table-result'><thead style='background-color: white'><tr><th>text</th><th>label</th><th>id</th></tr></thead><tbody><tr><td>a charming , banter-filled comedy . . . one of those airy cinematic bon bons whose aims -- and by extension , accomplishments -- seem deceptively slight on the surface .</td><td>1</td><td>60129542207</td></tr><tr><td>the attempt to build up a pressure cooker of horrified awe emerges from the simple fact that the movie has virtually nothing to show .</td><td>0</td><td>60129542175</td></tr><tr><td>marinated in clichés and mawkish dialogue .</td><td>0</td><td>60129542187</td></tr><tr><td>a dark comedy that goes for sick and demented humor simply to do so . the movie is without intent .</td><td>0</td><td>60129542230</td></tr><tr><td>the world needs more filmmakers with passionate enthusiasms like martin scorsese . but it doesn't need gangs of new york .</td><td>0</td><td>60129542260</td></tr><tr><td>as his circle of friends keeps getting smaller one of the characters in long time dead says 'i'm telling you , this is f * * * ed' . maybe he was reading the minds of the audience .</td><td>0</td><td>60129542214</td></tr><tr><td>scherfig's light-hearted profile of emotional desperation is achingly honest and delightfully cheeky .</td><td>1</td><td>60129542263</td></tr><tr><td>the plot is romantic comedy boilerplate from start to finish .</td><td>0</td><td>60129542188</td></tr><tr><td>humor in i spy is so anemic .</td><td>0</td><td>60129542206</td></tr><tr><td>the increasingly diverse french director has created a film that one can honestly describe as looking , sounding and simply feeling like no other film in recent history .</td><td>1</td><td>60129542242</td></tr></tbody></table></div>"]}}],"execution_count":0},{"cell_type":"markdown","source":["Here we instantiate the pre-trained model from Hugging Face, but load it with the fine-tuned model weights."],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"4c23b1c1-da08-4626-a5c1-b0598c8d19b8"}}},{"cell_type":"code","source":["device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nmodel_type = 'microsoft/deberta-v3-base'\nmax_length = 128\nhidden_dropout_prob = 0.\nattention_probs_dropout_prob = 0.\nnum_labels = 2\n\nmodel_folder = '/dbfs/model_outputs'\ntokenizer = AutoTokenizer.from_pretrained(model_type)\nconfig = AutoConfig.from_pretrained(model_type, hidden_dropout_prob=hidden_dropout_prob,\n                                    attention_probs_dropout_prob=attention_probs_dropout_prob, num_labels=num_labels)\nmodel = AutoModelForSequenceClassification.from_pretrained(model_folder, config=config).to(device)"],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"fe3f86d2-38cc-4e4a-adc8-424e74218eee"}},"outputs":[{"output_type":"display_data","metadata":{"application/vnd.databricks.v1+output":{"datasetInfos":[],"data":"Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.9/site-packages/transformers/convert_slow_tokenizer.py:434: UserWarning: The sentencepiece tokenizer that you are converting to a fast tokenizer uses the byte fallback option which is not implemented in the fast tokenizers. In practice this means that the fast version of the tokenizer can produce unknown tokens whereas the sentencepiece version would have converted these unknown tokens into a sequence of byte tokens matching the original piece of text.\n  warnings.warn(\nSpecial tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n","removedWidgets":[],"addedWidgets":{},"metadata":{},"type":"ansi","arguments":{}}},"output_type":"display_data","data":{"text/plain":["Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n/local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.9/site-packages/transformers/convert_slow_tokenizer.py:434: UserWarning: The sentencepiece tokenizer that you are converting to a fast tokenizer uses the byte fallback option which is not implemented in the fast tokenizers. In practice this means that the fast version of the tokenizer can produce unknown tokens whereas the sentencepiece version would have converted these unknown tokens into a sequence of byte tokens matching the original piece of text.\n  warnings.warn(\nSpecial tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"]}}],"execution_count":0},{"cell_type":"markdown","source":["We then create a [*TextClassificationPipeline*](https://huggingface.co/docs/transformers/v4.21.3/en/main_classes/pipelines#transformers.TextClassificationPipeline) object, which allows us to make model predictions directly from the input text. In this case, the *tokenization* is performed on-the-fly."],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"61ed5343-c85f-46de-a817-3f9fcc950d58"}}},{"cell_type":"code","source":["pipe = pipeline('text-classification', tokenizer=tokenizer, model=model, device=0)"],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"550d46b9-23d4-4a9c-b4b3-3e232900e621"}},"outputs":[],"execution_count":0},{"cell_type":"markdown","source":["Here is where we define the function that performs model inference over the testing data, using the *pipeline* object defined above.\n\nThe model inference is performed in parallel, over the Spark DataFrame, using Spark's [*Pandas UDF*](https://docs.microsoft.com/en-us/azure/databricks/spark/latest/spark-sql/udf-python-pandas) functionality."],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"cc52512c-d178-4488-9dd9-19cd1603ebc6"}}},{"cell_type":"code","source":["schema = StructType([\n  StructField('label', IntegerType()),\n  StructField('score', FloatType())\n])\n\n@pandas_udf(schema)\ndef predict(text: pd.Series) -> pd.DataFrame:\n    pred = pipe(text.to_list())\n    res = pd.DataFrame(pred)\n    res['label'][res['label'] == 'LABEL_1'] = 1\n    res['label'][res['label'] == 'LABEL_0'] = 0\n    res['label'] = res['label'].astype(int)\n    return res"],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"923df268-8bed-4297-ab87-24854c3fea63"}},"outputs":[],"execution_count":0},{"cell_type":"markdown","source":["The model inference function defined above is then executed through the *select* function from the *DataFrame* API. It returns a new column with the predicted labels and corresponding scores, given by the *pipeline* execution.\n\nWe then split the returned column into three separate columns corresponding to its components."],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"6358a9d9-ec57-4486-85ed-2f9570ca173f"}}},{"cell_type":"code","source":["sdf_test = sdf_test.select('id', 'text', 'label', predict('text').alias('predictions'))\n\nsdf_test = sdf_test.withColumns({'predicted_label': sdf_test['predictions'].getItem('label'),\n                                 'score': sdf_test['predictions'].getItem('score')}).drop('predictions')"],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"dbcb7f85-5213-4778-8ea9-164ef956205a"}},"outputs":[],"execution_count":0},{"cell_type":"markdown","source":["Display the testing data with the predicted values."],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"f3df26f6-3ae6-4b18-908b-eeceb9cc7717"}}},{"cell_type":"code","source":["start_time = time.time()\ndisplay(sdf_test.select(sdf_test.columns).limit(10))\nend_time = time.time()\n\nprint('\\nTotal Duration: %f' % (round(end_time-start_time, 3)))"],"metadata":{"application/vnd.databricks.v1+cell":{"title":"","showTitle":false,"inputWidgets":{},"nuid":"4ff31d6e-522c-46c6-9eaa-a935338d6e4b"}},"outputs":[{"output_type":"display_data","metadata":{"application/vnd.databricks.v1+output":{"overflow":false,"datasetInfos":[],"data":[[51539607677,"the ethos of the chelsea hotel may shape hawke's artistic aspirations , but he hasn't yet coordinated his own dv poetry with the beat he hears in his soul .",0,0,0.9518424],[51539607570,"a preposterously melodramatic paean to gang-member teens in brooklyn circa 1958 .",0,0,0.98857343],[51539607658,"director roger kumble offers just enough sweet and traditional romantic comedy to counter the crudity . and there's the inimitable diaz , holding it all together .",1,1,0.99498653],[51539607668,"weighty and ponderous but every bit as filling as the treat of the title .",1,1,0.99423236],[51539607653,"a reworking of die hard and cliffhanger but it's nowhere near as exciting as either .",0,0,0.9904231],[51539607640,"if i could have looked into my future and saw how bad this movie was , i would go back and choose to skip it . fortunately , you still have that option .",0,0,0.92581624],[51539607656,"a small gem from belgium .",1,1,0.98947996],[51539607565,"no surprises .",0,1,0.52954006],[51539607669,"the wild thornberrys movie is pleasant enough and the message of our close ties with animals can certainly not be emphasized enough .",1,1,0.9916142],[51539607638,"you don't know whether to admire the film's stately nature and call it classicism or be exasperated by a noticeable lack of pace . or both .",1,0,0.934046]],"plotOptions":{"displayType":"table","customPlotOptions":{},"pivotColumns":null,"pivotAggregation":null,"xColumns":null,"yColumns":null},"columnCustomDisplayInfos":{},"aggType":"","isJsonSchema":true,"removedWidgets":[],"aggSchema":[],"schema":[{"name":"id","type":"\"long\"","metadata":"{}"},{"name":"text","type":"\"string\"","metadata":"{}"},{"name":"label","type":"\"long\"","metadata":"{}"},{"name":"predicted_label","type":"\"integer\"","metadata":"{}"},{"name":"score","type":"\"float\"","metadata":"{}"}],"aggError":"","aggData":[],"addedWidgets":{},"metadata":{},"dbfsResultPath":null,"type":"table","aggOverflow":false,"aggSeriesLimitReached":false,"arguments":{}}},"output_type":"display_data","data":{"text/html":["<style scoped>\n  .table-result-container {\n    max-height: 300px;\n    overflow: auto;\n  }\n  table, th, td {\n    border: 1px solid black;\n    border-collapse: collapse;\n  }\n  th, td {\n    padding: 5px;\n  }\n  th {\n    text-align: left;\n  }\n</style><div class='table-result-container'><table class='table-result'><thead style='background-color: white'><tr><th>id</th><th>text</th><th>label</th><th>predicted_label</th><th>score</th></tr></thead><tbody><tr><td>51539607677</td><td>the ethos of the chelsea hotel may shape hawke's artistic aspirations , but he hasn't yet coordinated his own dv poetry with the beat he hears in his soul .</td><td>0</td><td>0</td><td>0.9518424</td></tr><tr><td>51539607570</td><td>a preposterously melodramatic paean to gang-member teens in brooklyn circa 1958 .</td><td>0</td><td>0</td><td>0.98857343</td></tr><tr><td>51539607658</td><td>director roger kumble offers just enough sweet and traditional romantic comedy to counter the crudity . and there's the inimitable diaz , holding it all together .</td><td>1</td><td>1</td><td>0.99498653</td></tr><tr><td>51539607668</td><td>weighty and ponderous but every bit as filling as the treat of the title .</td><td>1</td><td>1</td><td>0.99423236</td></tr><tr><td>51539607653</td><td>a reworking of die hard and cliffhanger but it's nowhere near as exciting as either .</td><td>0</td><td>0</td><td>0.9904231</td></tr><tr><td>51539607640</td><td>if i could have looked into my future and saw how bad this movie was , i would go back and choose to skip it . fortunately , you still have that option .</td><td>0</td><td>0</td><td>0.92581624</td></tr><tr><td>51539607656</td><td>a small gem from belgium .</td><td>1</td><td>1</td><td>0.98947996</td></tr><tr><td>51539607565</td><td>no surprises .</td><td>0</td><td>1</td><td>0.52954006</td></tr><tr><td>51539607669</td><td>the wild thornberrys movie is pleasant enough and the message of our close ties with animals can certainly not be emphasized enough .</td><td>1</td><td>1</td><td>0.9916142</td></tr><tr><td>51539607638</td><td>you don't know whether to admire the film's stately nature and call it classicism or be exasperated by a noticeable lack of pace . or both .</td><td>1</td><td>0</td><td>0.934046</td></tr></tbody></table></div>"]}},{"output_type":"display_data","metadata":{"application/vnd.databricks.v1+output":{"datasetInfos":[],"data":"\nTotal Duration: 24.730000\n","removedWidgets":[],"addedWidgets":{},"metadata":{},"type":"ansi","arguments":{}}},"output_type":"display_data","data":{"text/plain":["\nTotal Duration: 24.730000\n"]}}],"execution_count":0}],"metadata":{"application/vnd.databricks.v1+notebook":{"notebookName":"model_inference_pudf","dashboards":[],"notebookMetadata":{"pythonIndentUnit":4},"language":"python","widgets":{},"notebookOrigID":1267397042508521}},"nbformat":4,"nbformat_minor":0}