src/utils.py

import ast
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from dataclasses import dataclass

# Configs
class Config:
    code_gen_idt : str = "lvwerra/codeparrot-small"
    code_gen_config : dict = {
        "max_length" : 256
    }

class GenCode:
    def __init__(self,model_idt:str,gen_kwargs:dict) -> None:
        self.generation_pipe = pipeline("text-generation", model=model_idt)
        self.gen_kwargs = gen_kwargs
    def __call__(self, input : str):
        return self.generation_pipe(input, **self.gen_kwargs)[0]["generated_text"]
        


def return_lexer_map(code_snippet:str):
    return None


def return_parse_tree(code_snippet:str):
    return ast.dump(ast.parse(code_snippet))