import gradio as gr
from transformers import BertForQuestionAnswering
from transformers import BertTokenizerFast
import torch
from nltk.tokenize import word_tokenize
import timm

tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
model = timm.create_model('hf_hub:pseudolab/AI_Tutor_BERT', pretrained=True)
#model = BertForQuestionAnswering.from_pretrained("bert-base-uncased")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

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