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()