File size: 740 Bytes
52ab3ac 636aa45 52ab3ac 636aa45 52ab3ac 636aa45 52ab3ac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | 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") |