from flask import Flask, request, render_template
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

app = Flask(__name__)

# Load the fine-tuned model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("./finetune_model")
model = AutoModelForSeq2SeqLM.from_pretrained("./finetune_model")

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

@app.route('/summarize', methods=['POST'])
def summarize():
    article = request.form['article']
    inputs = tokenizer.encode(article, return_tensors="pt", truncation=True)
    summary_ids = model.generate(inputs, max_length=50, min_length=15, length_penalty=2.0, num_beams=4, early_stopping=True)
    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return render_template('index.html', summary=summary)

if __name__ == '__main__':
    import os
    if os.environ.get('FLASK_ENV') == 'development':
        app.run(debug=True, use_reloader=True)
    else:
        app.run(debug=False, use_reloader=False)