import pickle, torch, gradio as gr
from model.fih import Fih

with open("tokenizer.pkl", 'rb') as f:
    tokenizer = pickle.load(f)

model = Fih()

def respond(message, history):
    tokens = tokenizer.encode(message, return_tensor="pt")
    generated = tokens.tolist()
    for _ in range(60):
        logits = model.forward_pass(torch.tensor(generated))
        probs = torch.softmax(logits[-1] / 1.2, dim=-1)
        next_token = torch.multinomial(probs, 1).item()
        generated.append(next_token)

        if next_token == 0:
            break

    return tokenizer.decode(generated)

app = gr.ChatInterface(fn=respond, title="FIH1 | Chat 🐟")

# here be dragons...

app.launch()