import os
os.system("pip install git+https://github.com/openai/whisper.git")

import whisper
from flask import Flask, jsonify, request
import requests
import time
from transformers import pipeline


model = whisper.load_model("small")
pipe = pipeline(
    "automatic-speech-recognition",
    model="openai/whisper-small.en",
    chunk_length_s=15,
    device=model.device,
)

app = Flask(__name__)
app.config['TIMEOUT'] = 60 * 10 # 10 mins

@app.route("/")
def indexApi():
    return jsonify({"output": "okay"})

@app.route("/run", methods=['POST'])
def runApi():
    start_time = time.time()

    audio_url = request.form.get("audio_url")

    response = requests.get(audio_url)

    if response.status_code == requests.codes.ok:
        with open("audio.mp3", "wb") as f:
            f.write(response.content)
      
    else:
        return jsonify({
            "result": "Unable to save file, status code:  {response.status_code}" ,
        }), 400

    audio = "audio.mp3"

    audioOri = whisper.load_audio(audio)
    prediction = pipe(audioOri)["text"]

    end_time = time.time()
    total_time = end_time - start_time

    return jsonify({
        "audio_url": audio_url,
        "result": prediction,
        "exec_time_sec": total_time
    })

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)