Spaces:
Runtime error
Runtime error
import gradio as gr | |
from load_model import build | |
def custom_standardization(s): | |
s = tf.strings.lower(s) | |
s = tf.strings.regex_replace(s, f'[{re.escape(string.punctuation)}]', '') | |
s = tf.strings.join(['[START]', s, '[END]'], separator=' ') | |
return s | |
model = build() | |
def simple_gen(image, temperature=1): | |
initial = model.word_to_index([['[START]']]) # (batch, sequence) | |
img_features = model.feature_extractor(image[tf.newaxis, ...]) | |
tokens = initial # (batch, sequence) | |
for n in range(50): | |
preds = model((img_features, tokens)).numpy() # (batch, sequence, vocab) | |
preds = preds[:,-1, :] #(batch, vocab) | |
if temperature==0: | |
next = tf.argmax(preds, axis=-1)[:, tf.newaxis] # (batch, 1) | |
else: | |
next = tf.random.categorical(preds/temperature, num_samples=1) # (batch, 1) | |
tokens = tf.concat([tokens, next], axis=1) # (batch, sequence) | |
if next[0] == model.word_to_index('[END]'): | |
break | |
words = model.index_to_word(tokens[0, 1:-1]) | |
result = tf.strings.reduce_join(words, axis=-1, separator=' ') | |
return result.numpy().decode() | |
def transcribes(image): | |
result = [] | |
for t in [0,0.5,1]: | |
result.append(simple_gen(image, t)) | |
return result | |
gr.interface(fn=transcribes, | |
inputs=gr.Image(type="pil"), | |
outputs=["text","text","text"] | |
).launch() | |