codegen_api.py

from flask import jsonify, send_file, request
from main import *
import io, base64

def generate_code(prompt, output_path="output_code.py"):
    if codegen_model is None or codegen_tokenizer is None: return {"error": "Code generation model or tokenizer not initialized."}
    input_ids = codegen_tokenizer(prompt, return_tensors='pt').to(device); output = codegen_model.generate(input_ids, max_length=2048, temperature=0.7, top_p=0.9)
    code = codegen_tokenizer.decode(output[0], skip_special_tokens=True); return {"code": code}

def codegen_api(prompt):
    output = generate_code(prompt)
    if "error" in output: return {"error": output["error"]}
    code_base64 = base64.b64encode(output['code'].encode('utf-8')).decode('utf-8'); return {"code_base64": code_base64, "mimetype": "text/x-python", "filename": "output.py"}