From 9c14ee0cdd2e740700d8a330e8f4e5db6adaf357 Mon Sep 17 00:00:00 2001
From: eff-kay <faizan10114@gmail.com>
Date: Mon, 24 Feb 2025 04:21:30 -0500
Subject: [PATCH] query param

---
 example_configs/bert-eqa-query.yaml | 18 ++++++++++++++++++
 magemaker/runner.py                 |  8 ++++++++
 2 files changed, 26 insertions(+)
 create mode 100644 example_configs/bert-eqa-query.yaml

diff --git a/example_configs/bert-eqa-query.yaml b/example_configs/bert-eqa-query.yaml
new file mode 100644
index 0000000..b23dd86
--- /dev/null
+++ b/example_configs/bert-eqa-query.yaml
@@ -0,0 +1,18 @@
+deployment: !Deployment
+  destination: aws
+  # Endpoint name matches model_id for querying atm.
+  endpoint_name: huggingface-eqa-bert-base-cased
+  instance_count: 1
+  instance_type: ml.m5.xlarge
+
+models:
+- !Model
+  # Base model id that was finetuned
+  id: huggingface-eqa-bert-base-cased
+  source: custom
+
+  # local or S3 path
+  location: ./models/model.tar.gz
+query: !Query
+  # Query to send to the model
+  text: "What is the capital of France?"
\ No newline at end of file
diff --git a/magemaker/runner.py b/magemaker/runner.py
index 6de523e..6d8cc2e 100644
--- a/magemaker/runner.py
+++ b/magemaker/runner.py
@@ -161,11 +161,19 @@ def runner():
         help="Show magemaker version and exit"
     )
 
+    parser.add_argument(
+        '--query',
+        action='query',
+        help="path to YAML query configuration file",
+        type=str
+    )
+
     parser.add_argument(
         "--hf",
         help="Deploy a Hugging Face Model.",
         type=str
     )
+
     parser.add_argument(
         "--instance",
         help="EC2 instance type to deploy to.",