Create model.py

connections/model.py

@@ -0,0 +1,31 @@
+
+import transformers
+import torch
+
+
+class InferencePipeline:
+    def __init__(self, conf, api_key):
+        self.conf = conf
+        self.token = api_key
+        self.pipeline = self.get_model()
+
+    def get_model(self):
+
+        pipeline = transformers.pipeline(
+            "text-generation",
+            model=conf["model"]["model_name"],
+            model_kwargs={"torch_dtype": torch.bfloat16},
+            device_map=conf["model"]["device_map"],
+            token=self.token
+        )
+        
+        return pipeline
+
+    def infer(self, prompt):
+
+        outputs = pipeline(
+            prompt,
+            max_new_tokens=conf["model"]["max_new_tokens"],
+        )
+
+        return outputs[0]["generated_text"][-1]