# -*- coding: utf-8 -*-
"""Job-coach-testing-site-zeroshot.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/github/vanderbilt-data-science/job-coach-question-answering/blob/111-create-a-gradio-hf-space-to-test-whether-specific-information-is-in-a-coaching-document/job-coach-testing-zeroshot.ipynb
"""

import gradio as gr

import torch
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
from transformers import pipeline

classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
question_answerer = pipeline("question-answering", model = "bert-large-uncased-whole-word-masking-finetuned-squad")
def zshot(context, queries):
    answered=[]
    res=[]
    notAnswered=[]
    queries=queries.split("?")
    queries.pop(-1)
    for query in queries:
      query.strip()
      result = question_answerer(question = query, context=context)
      answered.append([query,result['answer']])
    
    for item in answered:
      result = ([text, classifier(text, item[0])])
      if result[1]['scores'][0] > 0.01:
        res.append(item[0] +"?   Answer: " + item[1])
      else:
        notAnswered.append(item[0])
    result1=''',
    '''.join(res) 
    result1='''Information Included in the Document: 
    '''+ result1
    result2=''',
    '''.join(notAnswered) 
    result2='''
    Information not included in the document, on no clearly stated: 
    '''+ result2

    return result1 + result2
    

app = gr.Interface(
    zshot,
    inputs=[
        gr.Textbox(label="Context", value=""),
        gr.Textbox(label="Queries", value=""),
    ],
    outputs=["text"],
)
app.launch()

"""#### Examples"""