Spaces:
Paused
Paused
| import gradio as gr | |
| from transformers import BertForQuestionAnswering | |
| from transformers import BertTokenizerFast | |
| import torch | |
| from nltk.tokenize import word_tokenize | |
| import timm | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased') | |
| #model = BertForQuestionAnswering.from_pretrained("bert-base-uncased") | |
| model = BertForQuestionAnswering.from_pretrained("CountingMstar/ai-tutor-bert-model").to(device) | |
| def get_prediction(context, question): | |
| inputs = tokenizer.encode_plus(question, context, return_tensors='pt').to(device) | |
| outputs = model(**inputs) | |
| answer_start = torch.argmax(outputs[0]) | |
| answer_end = torch.argmax(outputs[1]) + 1 | |
| answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs['input_ids'][0][answer_start:answer_end])) | |
| return answer | |
| def question_answer(context, question): | |
| prediction = get_prediction(context,question) | |
| return prediction | |
| def split(text): | |
| context, question = '', '' | |
| act = False | |
| tmp = '' | |
| for t in text: | |
| tmp += t | |
| if len(tmp) == 4: | |
| tmp = tmp[1:] | |
| if tmp == '///': | |
| act = True | |
| if act == True: | |
| question += t | |
| if act == False: | |
| context += t | |
| return context[:-2], question[1:] | |
| def greet(texts): | |
| context, question = split(texts) | |
| answer = question_answer(context, question) | |
| return answer | |
| # def greet(text): | |
| # context, question = split(text) | |
| # # answer = question_answer(context, question) | |
| # return context | |
| iface = gr.Interface(fn=greet, inputs="text", outputs="text") | |
| iface.launch() |