✨ implement the back-end call codegen and return the result

Author: chenqianhe

codegen_paddle/codegen.py

@@ -0,0 +1,39 @@
+import time
+import paddle
+from paddlenlp.utils.log import logger
+from paddlenlp.transformers import CodeGenTokenizer, CodeGenForCausalLM
+
+from .config import DefaultConfig
+
+generate_config = DefaultConfig()
+paddle.set_device(generate_config.device)
+paddle.set_default_dtype(generate_config.default_dtype)
+
+tokenizer = CodeGenTokenizer.from_pretrained(generate_config.model_name_or_path)
+model = CodeGenForCausalLM.from_pretrained(
+    generate_config.model_name_or_path,
+    load_state_as_np=generate_config.load_state_as_np)
+
+
+def gen_code(prompt: str) -> str:
+    start_time = time.time()
+    logger.info("Start generating code")
+    tokenized = tokenizer(prompt,
+                          truncation=True,
+                          return_tensors='pd')
+    output, _ = model.generate(
+        tokenized["input_ids"],
+        max_length=16,
+        min_length=generate_config.min_length,
+        decode_strategy=generate_config.decode_strategy,
+        top_k=generate_config.top_k,
+        repetition_penalty=generate_config.repetition_penalty,
+        temperature=generate_config.temperature,
+        use_faster=generate_config.use_faster,
+        use_fp16_decoding=generate_config.use_fp16_decoding)
+    logger.info("Finish generating code")
+    end_time = time.time()
+    logger.info(f"Time cost: {end_time - start_time}")
+    output = tokenizer.decode(output[0], skip_special_tokens=True)
+    logger.info(f"Generated code: {output}")
+    return output

codegen_paddle/config.py

@@ -0,0 +1,14 @@
+class DefaultConfig:
+    model_name_or_path = "Salesforce/codegen-350M-mono"
+    device = "cpu"
+    temperature = 0.5
+    top_k = 10
+    top_p = 1.0
+    repetition_penalty = 1.0
+    min_length = 0
+    max_length = 16
+    decode_strategy = "greedy_search"
+    load_state_as_np = True
+    use_faster = False
+    use_fp16_decoding = False
+    default_dtype = "float16" if use_faster and use_fp16_decoding else "float32"

codegen_paddle/handlers.py

@@ -7,6 +7,8 @@
 import tornado
 from tornado.web import StaticFileHandler
 
+from .codegen import gen_code
+
 
 class RouteHandler(APIHandler):
     # The following decorator should be present on all verb methods (head, get, post,
@@ -19,7 +21,9 @@ def get(self):
     @tornado.web.authenticated
     def post(self):
         input_data = self.get_json_body()
-        data = {"received": "{}".format(input_data["code"])}
+        prompt = input_data["code"]
+        res = gen_code(prompt)
+        data = {"received": "{}".format(res)}
         self.finish(json.dumps(data))
 
 
@@ -45,4 +49,3 @@ def setup_handlers(web_app, url_path):
     )
     handlers = [("{}/(.*)".format(doc_url), StaticFileHandler, {"path": doc_dir})]
     web_app.add_handlers(".*$", handlers)
-