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(new_tokens, ret_seq, txt):
    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 = int(new_tokens), num_return_sequences=int(ret_seq))
    return out


demo = gr.Interface(
    sentence_builder,
    [
        gr.Dropdown(
            ["10", "20", "30", "40", "50", "60", "80", "100"],
            label="Number of New Tokens",
            info="Choose how many tokens are required in output.",
            value="30"
        ),
        gr.Dropdown(
            ["1", "2", "3", "4", "5", "6", "8", "10"],
            label="Number of Return Sequences",
            info="Choose how many return sequences are required in output.",
            value="5"
        ),
        gr.Textbox("", label = "Input",
                   info="Give 8 words atleast, not to get concatenated with default words to make up it's minimum requirement."
                          )
    ],
    [
        gr.Textbox("", label = "Output")
    ],
    title="Shakespeare Drama Dialogue Mimick by GPT3",
    examples=[
               ["30", "5", "A horse! a horse! my kingdom for a horse!"],
               ["30", "5", "What’s In A Name? A Rose By Any Name Would Smell As Sweet"],
               ["30", "5", "If Music Be The Food Of Love, Play On"],
               ["30", "5", "Our Doubts Are Traitors And Make Us Lose The Good We Oft Might Win By Fearing To Attempt."],
               ["30", "5", "All That Glitters Is Not Gold."],
               ["30", "5", "Love All, Trust A Few, Do Wrong To None."],
               ["30", "5", "Love Looks Not With The Eyes, But With The Mind; And Therefore Is Winged Cupid Painted Blind"]
             ]
)

demo.launch(debug=True)