import gradio as gr
import openai
from t2a import text_to_audio
import joblib
from sentence_transformers import SentenceTransformer
import numpy as np

reg = joblib.load('text_reg.joblib')
model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
finetune = "davinci:ft-personal:autodrummer-v4-2022-11-01-22-44-58"

def get_note_text(prompt):
    prompt = prompt + " ->"
    # get completion from finetune
    response = openai.Completion.create(
        engine=finetune,
        prompt=prompt,
        temperature=0.7,
        max_tokens=100,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0,
        stop=["###"]
    )
    return response.choices[0].text.strip()

def get_drummer_output(prompt, openai_api_key):
    openai.api_key = openai_api_key
    note_text = get_note_text(prompt)
    # note_text = note_text + " " + note_text
    # note_text = "k n k n k n k n s n h n k n s n k n k n k n k n k n k n h n k n n"
    prompt_enc = model.encode([prompt])
    bpm = int(reg.predict(prompt_enc)[0]) + 20
    print(bpm, "bpm", "notes are", note_text)
    audio = text_to_audio(note_text, bpm)
    # audio to numpy
    audio = np.array(audio.get_array_of_samples(), dtype=np.float32)
    return (96000, audio)

iface = gr.Interface(fn=get_drummer_output, inputs=["text", "text"], outputs="audio")
iface.launch()