Update app.py

app.py

@@ -1,28 +1,23 @@
-from flask import Flask, request, jsonify, render_template
-from transformers import pipeline
+from flask import Flask, request, render_template
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Initialize Flask app
 app = Flask(__name__)
 
-# Load the fine-tuned model
-model_path = "./finetune_model"
-headline_generator = pipeline("text2text-generation", model=model_path)
+# Load the fine-tuned model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("path/to/finetune_model")
+model = AutoModelForSeq2SeqLM.from_pretrained("path/to/finetune_model")
 
 @app.route('/')
-def index():
+def home():
     return render_template('index.html')
 
-@app.route('/predict', methods=['POST'])
-def predict():
-    article = request.form.get('article')
-    if not article:
-        return jsonify({"error": "No article provided"}), 400
-    try:
-        prediction = headline_generator(article, max_length=50, num_return_sequences=1)
-        headline = prediction[0]['generated_text']
-        return jsonify({"headline": headline})
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
+@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__":
-    app.run(debug=True)
+if __name__ == '__main__':
+    app.run(debug=True)