Update app.py(2)

app.py

@@ -15,16 +15,31 @@ summarizer("Warm-up input.", max_length=30, min_length=5, do_sample=False)
 @app.route("/summarize", methods=["POST"])
 def summarize():
     try:
-        start_time = time.time()
         data = request.get_json()
         text = data["text"]
-        summary = summarizer(text, max_length=150, min_length=30, do_sample=False)
-        end_time = time.time()
+        max_len = int(data.get("max_length", 150))  # Default to 150 if not provided
+        min_len = int(data.get("min_length", 30))   # Default to 30 if not provided
+
+        summary_output = summarizer(
+            text,
+            max_length=max_len,
+            min_length=min_len,
+            do_sample=False
+        )
+
+        summary = summary_output[0]["summary_text"]
+
+        # Token count of the summary
+        tokenized = tokenizer(summary, return_tensors="pt")
+        token_count = len(tokenized["input_ids"][0])
+
+        return jsonify({
+            "summary": summary,
+            "token_count": token_count
+        })
 
-        print(f"Summary generated in {end_time - start_time:.2f} seconds")
-        return jsonify({"summary": summary[0]["summary_text"]})
     except Exception as e:
         return jsonify({"error": str(e)}), 400
-
+        
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)