import torch
import RohanGivenCode
from RohanGivenCode import *


save1_or_load0 = 0 # 1 => Save; 0 => Load

device = 'cpu'
if torch.cuda.is_available():
    device = 'cuda'
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
    device = "mps"
print(f"using device: {device}")

# SEED
torch.manual_seed(1337)
if torch.cuda.is_available():
    torch.cuda.manual_seed(1337)

# STOP
num_return_sequences = 5
max_length = 30


import gradio as gr
def sentence_builder(txt, new_tokens):
    txt_len = len(txt.split())
    if(txt_len < 9): # To make up minumum requirement of 9 words
        txt += " My lord, I claim your gift, my due by promise"

    t_loader = DataLoaderLite(B = 8, T = 1, text_input = txt)
    out = infer_the_model(device, t_loader, save1_or_load0 = 0, max_length = new_tokens)
    return out


demo = gr.Interface(
    sentence_builder,
    [
        gr.Textbox("", label = "Input", info="Give 8 words atleast, not to get concatenated with default words to make up it's minimum requirement."),
        gr.Dropdown(
            ["100", "200", "300", "400", "500", "1000", "2000"],
            label="New Tokens",
            info="Choose how many tokens required in output.",
            value="100"
        )
    ],
    [
        gr.Textbox("", label = "Output")
    ],
    title="Shakespeare Drama Dialogue Mimick by GPT3"
)

demo.launch(debug=True)