import requests
import json
import gradio as gr
import os

REST_API_KEY = os.environ.get("REST_API_KEY")

def kogpt_api(prompt, max_tokens = 32, temperature = 1.0, top_p = 1.0, n = 1):
    r = requests.post(
        url='https://api.kakaobrain.com/v1/inference/kogpt/generation',
        headers = {
            'Authorization': 'KakaoAK ' + REST_API_KEY,
            'Content-Type': 'application/json'
        },
        json = {
            'prompt': prompt,
            'max_tokens': max_tokens,
            'temperature': temperature,
            'top_p': top_p,
            'n': n
        }
    )
    
    response = json.loads(r.content)
    return response

def greet(prompt, max_tokens, temperature, top_p, n):
    response = kogpt_api(
        prompt = prompt,
        max_tokens = int(max_tokens),
        temperature = float(temperature),
        top_p = float(top_p),
        n = int(n)
    )
    return response['generations'][0]['text']

iface = gr.Interface(
    fn=greet,
    inputs=[
        gr.Textbox(lines=2, placeholder="Input Prompt Here...", label="textbox"),
        gr.Number(value=32),
        gr.Slider(0.1, 1.0, step=0.1, value=0.5),
        gr.Slider(0.1, 1.0, step=0.1, value=0.5),
        gr.Slider(1, 16, step=0.1, value=1)],
    outputs=["text"]
    )

iface.launch()