Upload handler.py

handler.py

@@ -0,0 +1,40 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import os
+import subprocess
+
+
+# Manually install bitsandbytes
+try:
+    import bitsandbytes
+except ImportError:
+    subprocess.check_call([os.sys.executable, "-m", "pip", "install", "bitsandbytes==0.39.1"])
+subprocess.check_call([os.sys.executable, "-m", "pip", "install", "accelerate==0.20.0"])
+
+class ModelHandler:
+    def __init__(self):
+        self.model = None
+        self.tokenizer = None
+
+    def load_model(self):
+        # Load token as env var
+        model_id = "NiCETmtm/Llama3_kw_gen_new"
+        token = os.getenv("HF_API_TOKEN")
+        # Load model & tokenizer
+        self.model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token, from_tf=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
+
+    def predict(self, inputs):
+        tokens = self.tokenizer(inputs, return_tensors="pt")
+        with torch.no_grad():
+            outputs = self.model.generate(**tokens)
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+
+model_handler = ModelHandler()
+model_handler.load_model()
+
+def inference(event, context):
+    inputs = event["data"]
+    outputs = model_handler.predict(inputs)
+    return {"predictions": outputs}