Upload 3 files

codegen.py

@@ -0,0 +1,36 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+
+DEVICE = 'cpu'
+TOKENIZER = None
+MODEL = None
+
+
+def setup(model: str, setup_torch: bool = False):
+    global TOKENIZER, MODEL, DEVICE
+    if setup_torch:
+        try:
+            import torch
+            torch.set_default_tensor_type(torch.cuda.FloatTensor)
+            # Make sure that we're using CPU when GPU isn't available
+            DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+        except:
+            print("ERROR: Can't set default tensor type to FloatTensor")
+
+    TOKENIZER = AutoTokenizer.from_pretrained(model)
+    MODEL = AutoModelForCausalLM.from_pretrained(model)
+
+
+def generate(token: str) -> str:
+    """
+    Generate some code using the loaded model given some input.
+
+    :param token: The input that will be passed into the tokenizer for the model to generate an output with.
+    :return: The string output.
+    """
+    if TOKENIZER is None or MODEL is None:
+        raise Exception("Model and tokenizer has not been setup.")
+
+    inputs = TOKENIZER(token, return_tensors='pt').to(DEVICE)
+    sample = MODEL.generate(**inputs, max_length=128)
+    return TOKENIZER.decode(sample[0], truncate_before_pattern=[r"\n\n^#", "^'''", "\n\n\n"])

wsgi.py

@@ -0,0 +1,4 @@
+from rubber_duck.app import app
+
+if __name__ == '__main__':
+    app.run(host="0.0.0.0", port="5002")