File size: 1,408 Bytes
dbfe252 6471fde 0945e74 dbfe252 0945e74 dbfe252 0945e74 52d9b6e c9b5ac1 52d9b6e 0945e74 52d9b6e c9b5ac1 7dc7f34 3970371 62fd772 c9b5ac1 8f99a74 0945e74 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import gradio as gr
import torch
import pytorch_lightning as pl
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from transformers import (
MT5ForConditionalGeneration,
MT5TokenizerFast,
)
model = MT5ForConditionalGeneration.from_pretrained(
"minjibi/qa",
return_dict=True,
)
tokenizer = MT5TokenizerFast.from_pretrained(
"minjibi/qa"
)
def predict(text):
input_ids = tokenizer.encode(text, return_tensors="pt", add_special_tokens=True)
generated_ids = model.generate(
input_ids=input_ids,
num_beams=5,
max_length=1000,
repetition_penalty=3.0, #default = 2.5
length_penalty=1.0,
early_stopping=True,
top_p=50, #default 50
top_k=20, #default 20
num_return_sequences=3,
)
preds = [
tokenizer.decode(
g,
skip_special_tokens=True,
clean_up_tokenization_spaces=True,
)
for g in generated_ids
]
output = ['Question: '+ text.replace('A', 'Answer') for text in preds]
#final_str = '\n'.join([f"{i+1}. {s}" for i, s in enumerate(output)])
final_str = '\n'.join([f"{i+1}. Question: {s.split('Answer')[0].strip()}\n Answer{s.split('Answer')[1].strip()}" for i, s in enumerate(output)])
return final_str
iface = gr.Interface(fn=predict, inputs="text", outputs="text")
iface.launch() |