update fine-tuning and deployment steps

GoogleCloudPlatform · Feb 12, 2025 · f232c33 · f232c33
1 parent e7d0712
commit f232c33
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 6 deletions.
diff --git a/tutorials-and-examples/mlflow/finetune-gemma/deploy-gemma2/kserve/input.json b/tutorials-and-examples/mlflow/finetune-gemma/deploy-gemma2/kserve/input.json
@@ -4,7 +4,7 @@
       "name": "args",
       "shape": [1],
       "datatype": "BYTES",
-      "data": ["Question: What is the total number of attendees with age over 30 at kubecon \"eu\"? Context: CREATE TABLE attendees (name VARCHAR, age INTEGER, kubecon VARCHAR)\nAnswer:"]
+      "data": ["Question: What is the total number of attendees with age over 30 at kubecon \"eu\"?\nContext: CREATE TABLE attendees (name VARCHAR, age INTEGER, kubecon VARCHAR)\n\nAnswer:"]
     }
   ]
 }
diff --git a/tutorials-and-examples/mlflow/finetune-gemma/deploy-gemma2/simple/deploy.yaml b/tutorials-and-examples/mlflow/finetune-gemma/deploy-gemma2/simple/deploy.yaml
@@ -19,7 +19,6 @@ spec:
       containers:
       - name: gemma2-9b-finetuned
         image: us-docker.pkg.dev/<PROJECT_ID>/gemma-deployment/finetune-gemma-mlflow:1.0.0
-        command: ["/bin/bash", "-c", "--"]
         env:
         - name: MODEL_PATH
           value: "/data/artifacts/<MLFLOW_RUN_ID>/artifacts/model/"

diff --git a/tutorials-and-examples/mlflow/finetune-gemma/deploy-gemma2/simple/main.py b/tutorials-and-examples/mlflow/finetune-gemma/deploy-gemma2/simple/main.py
@@ -10,4 +10,9 @@
 
 @app.get("/predict")
 async def predict(message: str):
-    return model.predict(message)
+    return model(
+        message,
+        max_new_tokens=64,
+        return_full_text=False,
+        repetition_penalty=1.
+    )
diff --git a/tutorials-and-examples/mlflow/finetune-gemma/finetune/finetune.py b/tutorials-and-examples/mlflow/finetune-gemma/finetune/finetune.py
@@ -7,6 +7,7 @@
     AutoTokenizer,
     BitsAndBytesConfig,
     TrainingArguments,
+    pipeline,
 )
 from peft import LoraConfig, PeftModel
 
@@ -248,8 +249,15 @@ def transform(data):
     model = PeftModel.from_pretrained(base_model, new_model)
     model = model.merge_and_unload()
 
+    generation_pipeline = pipeline(
+        task="text-generation",
+        model=model,
+        tokenizer=tokenizer,
+    )
+
     mlflow.log_params(peft_config.to_dict())
     mlflow.transformers.log_model(
-        transformers_model={"model": model, "tokenizer": tokenizer},
+        transformers_model=generation_pipeline,
         artifact_path="model",  # This is a relative path to save model files within MLflow run
+        prompt_template="{prompt}"
     )
diff --git a/tutorials-and-examples/mlflow/finetune-gemma/finetune/finetune.yaml b/tutorials-and-examples/mlflow/finetune-gemma/finetune/finetune.yaml
@@ -45,9 +45,9 @@ spec:
         - name: GRADIENT_ACCUMULATION_STEPS
           value: "2"
         - name: DATASET_LIMIT
-          value: "1000"
+          value: "2000"
         - name: MAX_SEQ_LENGTH
-          value: "512"
+          value: "256"
         - name: LOGGING_STEPS
           value: "5"
         - name: HF_TOKEN