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)