from transformers import ElectraTokenizer, ElectraForQuestionAnswering, pipeline 
from pprint import pprint 
from konlpy.tag import Mecab 
from flask import Flask, request, jsonify
import re

tokenizer = ElectraTokenizer.from_pretrained("monologg/koelectra-base-v3-finetuned-korquad") 
model = ElectraForQuestionAnswering.from_pretrained("monologg/koelectra-base-v3-finetuned-korquad") 
model = pipeline("question-answering", tokenizer=tokenizer, model=model, device=0) 
mecab = Mecab() 
pattern = r'\([^])]*\)'

app = Flask(__name__)

@app.route("/") 
def main(): 
    return { "QA_CIRCULUS" : "V1.1_TEST_2211"} 

@app.route("/qa", methods=['POST']) 
def qa(): 
    question = request.json['q']
    context = request.json['c']
    result = model({ "question" : question,  "context" : context }) 
    #print(result) 
    answer = result["answer"] 

    if answer.find('(') > -1 and answer.find(')') < 0:
        if(answer.startswith('(')):
            answer.replace("(","")
        else:
            answer = answer + ")"
    if answer.endswith('의'): 
        answer = answer.replace("의","")       
    answer = re.sub(pattern=pattern, repl='', string=answer )

    list = mecab.pos(result["answer"]) 
    print(list) 
    for word in list: 
        print(word[1]) 
        #if word[1] in ["JX","JKB","JKO"]: #Josa #Adjective 
        #if word[1].startswith('J'):
        #if answer.endswith('의'):
        #    answer = answer.replace('의','')
        #answer = answer.replace('이다','')
        #answer = answer.replace('라는','')
        if word[1].startswith('JKO') or word[1].startswith('JKS') or word[1].startswith('JKB') or word[1].startswith('JX') or word[1].startswith('JC'): 
            answer = answer.replace(word[0],"")
        if word[1].startswith('VCP'): 
            answer = answer.replace(word[0],"") 
        if word[1].startswith('SS'): 
            answer = answer.replace(word[0],"")   
        if word[1].endswith('F'): 
            answer = answer.replace(word[0],"")                        
        #if answer.find('(') > -1 and answer.find(')') < 0:
        #    answer = answer + ")"
        #if answer.find(''  
    print(result)
    
    result["answer"] = answer 
    return result 

#if __name__ == "__main__": 
#  uvicorn.run("qa_server:app", host="0.0.0.0", port=8000) 
if __name__ == '__main__':
    app.run(host="0.0.0.0",port=5000)