Rishabh100/code-comment-generator

1# app.py2import os3from flask import Flask, request, jsonify, render_template4from transformers import RobertaTokenizer, T5ForConditionalGeneration5 6app = Flask(__name__)7 8# Use smaller model for faster deployment9MODEL_NAME = "Salesforce/codet5-small"10 11# Initialize variables12tokenizer = None13model = None14 15def load_model():16    """Lazily load model to avoid timeout issues"""17    global tokenizer, model18    if tokenizer is None or model is None:19        tokenizer = RobertaTokenizer.from_pretrained(MODEL_NAME)20        model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME)21    return tokenizer, model22 23def generate_comment(code_snippet):24    try:25        tokenizer, model = load_model()26        input_text = "Generate comment: " + code_snippet27        input_ids = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)28        29        output = model.generate(30            input_ids,31            max_length=128,32            num_beams=4,33            early_stopping=True34        )35        36        return tokenizer.decode(output[0], skip_special_tokens=True)37    except Exception as e:38        return f"Error: {str(e)}"39 40@app.route('/')41def index():42    return render_template('index.html')43 44@app.route('/generate', methods=['POST'])45def generate():46    code = request.json.get('code', '')47    if not code.strip():48        return jsonify({'comment': 'Please enter code'})49    50    comment = generate_comment(code)51    return jsonify({'comment': comment})52 53if __name__ == '__main__':54    app.run(debug=True, host='0.0.0.0', port=int(os.environ.get('PORT', 8080)))55