| 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") |
