微調擁抱臉模型一個GPU

從變形金剛進口AutoTokenizer分詞器=AutoTokenizer。from_pretrained(base_model)deftokenize_function(例子):返回分詞器(例子(“文本”),填充=假,截斷=真正的)train_test_tokenized=train_test_dataset。地圖(tokenize_function,批處理=真正的)

進口numpy作為np進口評估度規=評估。負載(“準確性”)defcompute_metrics(eval_pred):分對數,標簽=eval_pred預測=np。argmax(分對數,軸= -1)返回度規。計算(預測=預測,引用=標簽)

從變形金剛進口AutoModelForSequenceClassification模型=AutoModelForSequenceClassification。from_pretrained(base_model,num_labels=len(label2id),label2id=label2id,id2label=id2label)

從變形金剛進口TrainingArguments,教練training_args=TrainingArguments(output_dir=training_output_dir,evaluation_strategy=“時代”)

從變形金剛進口DataCollatorWithPaddingdata_collator=DataCollatorWithPadding(分詞器)

教練=教練(模型=模型,arg遊戲=training_args,train_dataset=train_test_dataset(“訓練”),eval_dataset=train_test_dataset(“測試”),compute_metrics=compute_metrics,data_collator=data_collator,)

從變形金剛進口管道與mlflow。start_run()作為運行:教練。火車()教練。save_model(model_output_dir)管=管道(“文本分類”,模型=AutoModelForSequenceClassification。from_pretrained(model_output_dir),batch_size=1,分詞器=分詞器)model_info=mlflow。變形金剛。log_model(transformers_model=管,artifact_path=“分類”,input_example=“大家好!”,)

model_info=mlflow。變形金剛。log_model(transformers_model={“模型”:教練。模型,“記號賦予器”:分詞器},任務=“文本分類”,artifact_path=“text_classifier”,input_example=(“MLflow很棒!”,“MLflow磚是令人敬畏的!”),)

logged_model=”:/{run_id}/{model_artifact_path}”。格式(run_id=運行。信息。run_id,model_artifact_path=model_artifact_path)#負荷模型作為引發UDF。覆蓋result_type如果模型不返回值的兩倍。loaded_model_udf=mlflow。pyfunc。spark_udf(火花,model_uri=logged_model,result_type=“字符串”)測試=測試。選擇(測試。文本,測試。標簽,loaded_model_udf(測試。文本)。別名(“預測”))顯示(測試)