| from flask import Flask, render_template, request, jsonify
|
| from transformers import PegasusTokenizer, PegasusForConditionalGeneration
|
| import torch
|
| import re
|
|
|
| app = Flask(__name__)
|
|
|
| device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
| MODEL_PATH = "TheMarvellousOne/pegasus-meeting-summarizer"
|
|
|
| tokenizer = None
|
| model = None
|
|
|
|
|
| def load_model():
|
| global tokenizer, model
|
|
|
| if model is None:
|
| print("Loading model...")
|
|
|
| tokenizer = PegasusTokenizer.from_pretrained("google/pegasus-xsum")
|
|
|
| model = PegasusForConditionalGeneration.from_pretrained(
|
| MODEL_PATH,
|
| subfolder="pegasus_mixed_finetuned"
|
| ).to(device)
|
|
|
| model.eval()
|
|
|
| print("Model loaded.")
|
|
|
| def clean_transcript(text):
|
| """Remove extra whitespace from transcript"""
|
| text = str(text)
|
| text = re.sub(r"\s+", " ", text)
|
| return text.strip()
|
|
|
|
|
| def summarize_pegasus(text, tokenizer, model, device="cuda"):
|
| text = clean_transcript(text)
|
|
|
| inputs = tokenizer(
|
| text,
|
| return_tensors="pt",
|
| max_length=512,
|
| truncation=True
|
| ).to(device)
|
|
|
| generate_kwargs = {
|
| "max_length": 256,
|
| "min_length": 50,
|
| "num_beams": 8,
|
| "length_penalty": 1.5,
|
| "repetition_penalty": 1.5,
|
| "no_repeat_ngram_size": 3,
|
| "early_stopping": True,
|
| "do_sample": True,
|
| "temperature": 0.7,
|
| "top_p": 0.9,
|
| }
|
|
|
| with torch.no_grad():
|
| summary_ids = model.generate(
|
| **inputs,
|
| **generate_kwargs
|
| )
|
|
|
| return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
|
|
|
|
|
| @app.route("/")
|
| def home():
|
| load_model()
|
| return render_template("index.html")
|
|
|
|
|
| @app.route("/summarize", methods=["POST"])
|
| def summarize():
|
| load_model()
|
| transcript = request.json["text"]
|
|
|
| summary = summarize_pegasus(
|
| transcript,
|
| tokenizer,
|
| model,
|
| device
|
| )
|
|
|
| return jsonify({
|
| "summary": summary
|
| })
|
|
|
| if __name__ == "__main__":
|
| app.run(host="0.0.0.0", port=7860) |
