codestyle-diffuser / src /inference.py
hmnshudhmn24's picture
Upload 30 files
b2edbd0 verified
from .code_parser import CodeParser
from .style_encoder import StyleEncoder
from .style_diffuser import StyleDiffuser
from .rewrite_engine import RewriteEngine
class CodeStylePipeline:
def __init__(self):
self.p=CodeParser(); self.e=StyleEncoder()
self.d=StyleDiffuser(); self.r=RewriteEngine()
def __call__(self,code,style):
lines=self.p.parse(code)
prof=self.e.encode(style)
lines=self.d.diffuse(lines,prof['indent'])
lines=self.r.rewrite(lines,prof['max_line_length'])
return '\n'.join(lines)