action deploy - add model directory

Author: gabriben

model_artifacts/inference.py

@@ -0,0 +1,30 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+def model_fn(model_dir, *args):
+    # Load model from HuggingFace Hub
+    bnb_config = BitsAndBytesConfig(
+      load_in_4bit=True,
+      bnb_4bit_quant_type="nf4",
+      bnb_4bit_use_double_quant=True,
+      bnb_4bit_compute_dtype=torch.bfloat16
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+      model_dir,
+      device_map="auto",
+      quantization_config=bnb_config
+  )
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    return model, tokenizer
+
+def predict_fn(data, model_and_tokenizer, *args):
+    # destruct model and tokenizer
+    model, tokenizer = model_and_tokenizer
+    # Tokenize sentences
+    sentences = data.pop("inputs", data)
+    tokenizer.padding_side = "left"
+    tokenizer.pad_token = tokenizer.eos_token
+    model.config.pad_token_id = model.config.eos_token_id
+    inputs = tokenizer(sentences, return_tensors="pt", padding=True).to(model.device)
+    output_sequences = model.generate(**inputs, max_new_tokens=20)
+    return tokenizer.batch_decode(output_sequences, skip_special_tokens=True)

model_artifacts/requirements.txt

@@ -0,0 +1,5 @@
+bitsandbytes==0.44.1
+accelerate==1.6.0
+transformers==4.51.1
+torch==2.5.0
+torchvision==0.20

src/wraval/actions/action_deploy.py

@@ -6,6 +6,8 @@
 import boto3
 import json
 
+MODEL_DIRECTORY = '../../model_artifacts'
+
 def cleanup_endpoints(endpoint_name):
 
     sagemaker_client = boto3.client("sagemaker", region_name='us-east-1')