Update app.py

app.py

@@ -1,47 +1,28 @@
-from flask import Flask, request, render_template
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
+from flask import Flask, request, jsonify, render_template
+from transformers import pipeline
 
 # Initialize Flask app
 app = Flask(__name__)
 
-# Load the fine-tuned model and tokenizer
-model_dir = "./finetune_model"
-tokenizer = AutoTokenizer.from_pretrained(model_dir)
-
-
-
-model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
-model.eval()  # Set model to evaluation mode
-
-# Generate headline from article
-def generate_headline(article, max_length=128, num_beams=5):
-    # Tokenize the input article
-    inputs = tokenizer(article, max_length=256, truncation=True, return_tensors="pt", padding="max_length")
-
-    # Move inputs to the same device as the model
-    input_ids = inputs['input_ids'].to(model.device)
-    attention_mask = inputs['attention_mask'].to(model.device)
-
-    # Generate headline
-    outputs = model.generate(input_ids, attention_mask=attention_mask, max_length=max_length, num_beams=num_beams, early_stopping=True)
-    headline = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    return headline
-
-# Home route to render the form and handle POST requests
-@app.route('/', methods=['GET', 'POST'])
-def home():
-    headline = None
-    if request.method == 'POST':
-        article = request.form.get('article')
-        if article:
-            headline = generate_headline(article)
-    return render_template('index.html', headline=headline)
-
-# Run the app
-if __name__ == '__main__':
-    app.run(app.run(port=5001, debug=False))
-
-
-
+# Load the fine-tuned model
+model_path = "./finetune_model"
+headline_generator = pipeline("text2text-generation", model=model_path)
+
+@app.route('/')
+def index():
+    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
+
+if __name__ == "__main__":
+    app.run(debug=True)