from transformers import T5ForConditionalGeneration, AutoTokenizer
import gradio as gr
import os
cwd = os.getcwd()

print("Starting")
device = "cpu"

checkpoint = "./models/codet5_large_typescript"

tokenizer = AutoTokenizer.from_pretrained(checkpoint, local_files_only=True)

print("Loading Model")
model = T5ForConditionalGeneration.from_pretrained(checkpoint, local_files_only=True).to(device)
print("Model Loaded")

def completeCode(prompt):
	inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
	outputs = model.generate(inputs, max_new_tokens = 100, early_stopping = True ).to(device)
	return tokenizer.decode(outputs[0])

iface = gr.Interface(fn=completeCode, inputs="text", outputs="text")
iface.launch()
print("Ready")