Update app.py

app.py

@@ -1,6 +1,14 @@
+from flask import Flask, request, jsonify
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 import gradio as gr
 from huggingface_hub import InferenceClient
 
+app = Flask(__name__)
+
+tokenizer = AutoTokenizer.from_pretrained("sshleifer/distilbart-cnn-12-6")
+model = AutoModelForSeq2SeqLM.from_pretrained("sshleifer/distilbart-cnn-12-6")
+
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 """
@@ -40,6 +48,18 @@ def respond(
         yield response
 
 
+@app.route('/summarize', methods=['POST'])
+def summarize():
+    data = request.json
+    text = data.get('text', '')
+    if not text:
+        return jsonify({'error': 'No text provided'}), 400
+
+    inputs = tokenizer([text], max_length=1024, return_tensors="pt", truncation=True)
+    summary_ids = model.generate(inputs["input_ids"], max_length=60, min_length=10, length_penalty=2.0, num_beams=4, early_stopping=True)
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return jsonify({'summary': summary})
+
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
 """
@@ -59,6 +79,7 @@ demo = gr.ChatInterface(
     ],
 )
 
-
-if __name__ == "__main__":
-    demo.launch()
+if __name__ == '__main__':
+    from waitress import serve
+    serve(app, host='0.0.0.0', port=5005)
+    demo.launch()