codegen: Stage the code

.gitignore

@@ -0,0 +1,28 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*$py.class
+
+# VS code files
+.vscode
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Logs
+logs.log
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/

README.md

@@ -1,10 +1,37 @@
----
-title: Codegen
-emoji: 📊
-colorFrom: indigo
-colorTo: yellow
-sdk: docker
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# UNISYS Codegen
+
+A FastAPI project use to generate the code snippets based on the input given.
+
+We are using Granite-3B-Code-Base code model designed for code generative tasks.
+
+
+## 🌞 How to start app
+Make sure you have python 3.10 or higher installed.
+First install the dependencies.
+
+Installation of transformers:
+```
+git clone https://github.com/huggingface/transformers
+cd transformers/
+pip install ./
+```
+
+Installing required packages:
+
+```bash
+pip install -r requirements.txt
+```
+
+### Manually
+
+- Start with unicorn
+```bash
+pip install unicorn
+uvicorn index:app --reload
+```
+
+### API documentation: (Change the doamin-port accordingly)
+
+```
+http://localhost:8000/docs/
+```

controllers/__init__.py

@@ -0,0 +1,6 @@
+import os
+from fastapi_router_controller import ControllerLoader
+
+this_dir = os.path.dirname(__file__)
+
+ControllerLoader.load(this_dir, __package__)

controllers/generate_controller.py

@@ -0,0 +1,53 @@
+import asyncio
+from fastapi_router_controller import Controller
+from fastapi import APIRouter, Depends, HTTPException
+from services.ibm_model.ibm_extract_code_block import IbmExtractCodeblock
+from services.ibm_model.ibm_text_generator import IbmTextGenerator
+from services.model_generator import ModelGenerator
+
+from utils.logger import Logger
+
+logger = Logger.get_logger(__name__)
+
+router = APIRouter(prefix='/v1')
+controller = Controller(router, openapi_tag={
+    'name': 'Generate the Code Snippets',
+})
+
+model = ModelGenerator()
+
+
+@controller.use()
+@controller.resource()
+class GenerateController():
+    def __init__(
+            self,
+            service: ModelGenerator = Depends()) -> None:
+        self.model_generator = service
+
+    @controller.route.get(
+        '/generate',
+        tags=['generate-code'],
+        summary='Generates the code for the given input')
+    async def generate_code(self, input: str):
+        try:
+            if not input:
+                logger.error('Input is required.')
+                raise HTTPException(
+                    status_code=500, detail='Input is required.')
+
+            ibm_generate_text_visitor = IbmTextGenerator()
+            generated_text = await self.model_generator.acceptTextGenerator(ibm_generate_text_visitor, input)
+
+            ibm_extract_code_block_visitor = IbmExtractCodeblock()
+            code_block = self.model_generator.acceptExtractCodeBlock(
+                ibm_extract_code_block_visitor, generated_text)
+
+            logger.info('Output: {}'.format(generated_text))
+
+            return {"data": code_block}
+        except asyncio.CancelledError:
+            logger.error(
+                'Canceling network request due to disconnect in client.')
+        except Exception as error:
+            logger.error('Error {}'.format(error))

index.py

@@ -0,0 +1,21 @@
+import controllers
+
+from fastapi import FastAPI
+from fastapi_router_controller import Controller, ControllersTags
+
+from utils.config import Config
+from utils.middleware import LogIncomingRequest
+from utils.middleware.request_cancellation import RequestCancellation
+
+#########################################
+#### Configure the main application #####
+#########################################
+app = FastAPI(
+    title='{}'.format(Config.read('app', 'name')),
+    openapi_tags=ControllersTags)
+
+app.add_middleware(LogIncomingRequest)
+app.add_middleware(RequestCancellation)
+
+for router in Controller.routers():
+    app.include_router(router)

requirements.txt

@@ -0,0 +1,2 @@
+fastapi
+fastapi-router-controller

services/ibm_model/ibm_extract_code_block.py

@@ -0,0 +1,37 @@
+import re
+from services.model_visitor import ModelVisitor
+
+
+class IbmExtractCodeblock(ModelVisitor):
+
+    def visit(self, _, data):
+        return self._get_code_block(data)
+
+    def _get_code_block(self, data):
+        r"""
+            Extracts text blocks from the input string based on a specific pattern.
+            Args:
+                data (str): The input string containing text blocks.
+            Returns:
+                str: A text block of output which contains code extracted from the input string.
+            Regex Pattern:
+            (?:### Output: ([\s\S]*?))(?:\<\|endoftext\|\>|\Z)|```(?:\w+)?\n(.*?)\n```
+            - (?:### Output: ([\s\S]*?)): This part matches patterns that start with '### Output:' 
+            followed by any characters including newlines, capturing them within a group.
+            - (?:\<\|endoftext\|\>|\Z): This part matches either the string <|endoftext|>
+            or the end of the string (\Z).
+            - |: This is an OR operator, meaning the regex will match either the pattern 
+            before or after it.
+            - ```(?:\w+)?\n(.*?)\n```: This part matches patterns enclosed within backticks (```), 
+            possibly preceded by one or more word characters (\w+), capturing any characters 
+            including newlines.
+        """
+        pattern = r'(?:### Output: ([\s\S]*?))(?:\<\|endoftext\|\>|\Z)|```(?:\w+)?\n(.*?)\n```'
+        matches = re.findall(pattern, data, re.DOTALL)
+        code = []
+        for match in matches:
+            if match[0]:
+                code.append(match[0].strip())
+            elif match[1]:
+                code.append(match[1].strip())
+        return ''.join(code)

services/ibm_model/ibm_text_generator.py

@@ -0,0 +1,65 @@
+import asyncio
+import torch
+from datetime import datetime
+from services.model_visitor import ModelVisitor
+from utils.logger import Logger
+
+logger = Logger.get_logger(__name__)
+
+
+class IbmTextGenerator(ModelVisitor):
+
+    async def visit(self, model_generator, input_text, max_length_per_chunk=50):
+        return await self._generate_text(model_generator, input_text, max_length_per_chunk)
+
+    async def _generate_text_chunk(self, model_generator, input_ids, max_length_per_chunk):
+        with torch.no_grad():
+            outputs = await asyncio.to_thread(model_generator.model.generate, input_ids, max_new_tokens=max_length_per_chunk)
+            continuation = model_generator.tokenizer.decode(
+                outputs[0], skip_special_tokens=False)
+
+            logger.info('Chunk generated: {}'.format(continuation))
+
+            return continuation
+
+    async def _generate_text(self, model_generator, input_text, max_length_per_chunk):
+        """
+            Generates the text based on input provided
+            Args:
+                input_text (str): The input string containing text blocks.
+                max_length_per_chunk: Max length per chunk (Default: 50 / Optional)
+        """
+        try:
+            start_time = datetime.now()
+
+            logger.info('Started at: {}'.format(
+                start_time.strftime(model_generator._format_data_time)))
+
+            input_ids = model_generator.tokenizer.encode(
+                input_text, return_tensors='pt').to(model_generator.device)
+
+            output_text = input_text
+            while True:
+                continuation = await self._generate_text_chunk(
+                    model_generator, input_ids, max_length_per_chunk)
+
+                new_text = continuation[len(model_generator.tokenizer.decode(
+                    input_ids[0], skip_special_tokens=False)):]
+                output_text += new_text
+                input_ids = model_generator.tokenizer.encode(
+                    output_text, return_tensors='pt').to(model_generator.device)
+                if "<|endoftext|>" in new_text or new_text.count('```') > 1:
+                    break
+
+            end_time = datetime.now()
+
+            logger.info('Output generated at: {}'.format(
+                end_time.strftime(model_generator._format_data_time)))
+
+            logger.info('Time taken: {}'.format(end_time - start_time))
+
+            return output_text
+        except asyncio.CancelledError:
+            logger.error(
+                'Cancelling model generation due to disconnection in network.')
+            return ""

services/model_generator.py

@@ -0,0 +1,47 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from utils.config import Config
+from utils.logger import Logger
+
+logger = Logger.get_logger(__name__)
+
+
+class ModelGenerator:
+    """
+    Singleton class responsible for generating text using a specified language model.
+
+    This class initializes a language model and tokenizer, and provides methods 
+    to generate text and extract code blocks from generated text.
+
+    Attributes:
+        device (torch.device): Device to run the model on (CPU or GPU).
+        model (AutoModelForCausalLM): Language model for text generation.
+        tokenizer (AutoTokenizer): Tokenizer corresponding to the language model.
+
+    Methods:
+        acceptTextGenerator(self, visitor, *args, **kwargs):
+            Accepts a visitor to generates text based on the input provided with the model generator.
+        acceptExtractCodeBlock(self, visitor, *args, **kwargs):
+            Accepts a visitor to extract code blocks from the output text.
+    """
+    _instance = None
+    _format_data_time = "%Y-%m-%d %H:%M:%S"
+
+    def __new__(cls, model_name=Config.read('app', 'model')):
+        if cls._instance is None:
+            cls._instance = super(ModelGenerator, cls).__new__(cls)
+            cls._instance._initialize(model_name)
+        return cls._instance
+
+    def _initialize(self, model_name):
+        self.device = torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu")
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name).to(self.device)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    def acceptTextGenerator(self, visitor, *args, **kwargs):
+        return visitor.visit(self, *args, **kwargs)
+
+    def acceptExtractCodeBlock(self, visitor, *args, **kwargs):
+        return visitor.visit(self, *args, **kwargs)

services/model_visitor.py

@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+
+
+class ModelVisitor(ABC):
+    """
+    Abstract base class for model visitors.
+
+    This class defines the interface for visiting a model generator. 
+    Subclasses must implement the visit method to define 
+    specific behaviors for different types of model generators.
+
+    Methods:
+        visit(generator, *args, **kwargs):
+            Abstract method to visit the model generator. Subclasses 
+            must override this method to provide specific functionality.
+
+    Example:
+        class IbmTextGenerator(ModelVisitor):
+            def visit(self, model_generator, *args, **kwargs):
+                # Implement specific behavior here
+                pass
+    """
+
+    @abstractmethod
+    def visit(self, generator, *args, **kwargs):
+        pass

utils/config/__init__.py

@@ -0,0 +1,15 @@
+import os
+from pathlib import Path
+from configparser import ConfigParser
+
+this_dir = Path(__file__).parent
+conf_dir = this_dir / 'properties.ini'
+
+parser = ConfigParser(os.environ)
+parser.read(conf_dir, encoding="utf8")
+
+
+class Config():
+    @staticmethod
+    def read(section, property, default=None):
+        return parser.get(section, property) or default

utils/config/properties.ini

@@ -0,0 +1,12 @@
+[app]
+name=Unisys Code Gen
+env=%(ENV)
+
+# we are currently using this model
+model=ibm-granite/granite-3b-code-base
+
+[log]
+level=INFO
+filename=./logs.log
+dateformat=%%Y-%%m-%%dT%%H:%%M:%%S
+format=%%(asctime)s.%%(msecs)03d %%(levelname)5s %%(name)s - %%(message)s

utils/logger.py

@@ -0,0 +1,38 @@
+import logging
+
+from utils.config import Config
+
+
+class Logger():
+    CONFIG_KEY = 'log'
+
+    @staticmethod
+    def get_level():
+        return Config.read(Logger.CONFIG_KEY, 'level')
+
+    @staticmethod
+    def get_filename():
+        return Config.read(Logger.CONFIG_KEY, 'filename')
+
+    @staticmethod
+    def get_format():
+        return Config.read(Logger.CONFIG_KEY, 'format')
+
+    @staticmethod
+    def get_date_format():
+        return Config.read(Logger.CONFIG_KEY, 'dateformat')
+
+    @staticmethod
+    def get_logger(name):
+        logger = logging.getLogger(name)
+        logger.setLevel(Logger.get_level())  # type: ignore
+
+        formatter = logging.Formatter(
+            Logger.get_format(),
+            Logger.get_date_format())
+
+        file_hdlr = logging.FileHandler(Logger.get_filename())  # type: ignore
+        file_hdlr.setFormatter(formatter)
+        logger.addHandler(hdlr=file_hdlr)
+
+        return logger

utils/middleware/__init__.py

@@ -0,0 +1,5 @@
+from utils.middleware.log_incoming_request import LogIncomingRequest
+
+__all__ = [
+    'LogIncomingRequest'
+]

utils/middleware/log_incoming_request.py

@@ -0,0 +1,32 @@
+import time
+from typing import List
+from fastapi import Request
+from fastapi.routing import APIRoute
+from starlette.middleware.base import BaseHTTPMiddleware
+from utils.logger import Logger
+
+logger = Logger.get_logger(__name__)
+
+
+class LogIncomingRequest(BaseHTTPMiddleware):
+    def __get_request_handler(_, req: Request):  # type: ignore
+        # get controller from request
+        routes: List[APIRoute] = req.app.routes
+        for route in routes:
+            if route.path_regex.match(req.url.path) and req.method in route.methods:
+                return route.endpoint.__name__ if hasattr(route.endpoint, '__name__') else 'fastapi_core'
+
+    async def dispatch(self, request: Request, call_next):
+        func_name = self.__get_request_handler(request)
+        request.state.func_name = func_name
+
+        logger.info('{} - start'.format(func_name))
+        start_time = time.time()
+
+        response = await call_next(request)
+
+        process_time = (time.time() - start_time) * 1000
+        formatted_process_time = '{0:.2f}'.format(process_time)
+        logger.info('{} - end in time (ms): {}'.format(func_name,
+                    formatted_process_time))
+        return response

utils/middleware/request_cancellation.py

@@ -0,0 +1,39 @@
+import asyncio
+from utils.logger import Logger
+
+logger = Logger.get_logger(__name__)
+
+
+class RequestCancellation:
+    """
+        RequestCancellation middleware handles request canceling
+        * In case of API routes where very frequent/expensive requests are made.
+    """
+
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        queue = asyncio.Queue()
+
+        async def message_poller(sentinel, handler_task):
+            nonlocal queue
+            while True:
+                message = await receive()
+                if message["type"] == "http.disconnect":
+                    handler_task.cancel()
+                    return sentinel
+                await queue.put(message)
+
+        sentinel = object()
+        handler_task = asyncio.create_task(self.app(scope, queue.get, send))
+        asyncio.create_task(message_poller(sentinel, handler_task))
+
+        try:
+            return await handler_task
+        except asyncio.CancelledError:
+            logger.info('Task Cancellatation Requested.')