from flask import Flask, render_template, request, jsonify
from transformers import T5ForConditionalGeneration, T5Tokenizer

app = Flask(__name__)

# Load the trained model
model = T5ForConditionalGeneration.from_pretrained("text2sql/model")
tokenizer = T5Tokenizer.from_pretrained("text2sql/model")

@app.route('/')
def home():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    if request.method == 'POST':
        input_text = "translate English to SQL : " + request.form['input_text']
        
        # Preprocess input
        input_ids = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)["input_ids"]
        
        # Generate SQL query
        output = model.generate(input_ids)
        predicted_sql = tokenizer.decode(output[0], skip_special_tokens=True)
        
        return render_template('index.html', input_text=input_text, predicted_sql=predicted_sql)

if __name__ == '__main__':
    app.run(debug=True, port=7860)