Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import AutoModelWithLMHead, AutoTokenizer | |
model_name = "distilbert-base-uncased" | |
model = AutoModelWithLMHead.from_pretrained(model_name) | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
def complete_sentence(sentence): | |
input_ids = tokenizer.encode(sentence, return_tensors="pt") | |
output = model.generate(input_ids, max_length=50, num_return_sequences=1) | |
completed_sentence = tokenizer.decode(output[0], skip_special_tokens=True) | |
return completed_sentence | |
iface = gr.Interface(fn=complete_sentence, inputs="text", outputs="text") | |
if __name__ == "__main__": | |
iface.launch() | |