import streamlit as st
from sentence_transformers import SentenceTransformer, util
from transformers import (AutoModelForQuestionAnswering,
                          AutoTokenizer, pipeline)

import pandas as pd
import regex as re

# Select model for question answering
model_name = "deepset/roberta-base-squad2"

# Load model & tokenizer
model = AutoModelForQuestionAnswering.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Create pipeline
pipe = pipeline('question-answering', model=model_name, tokenizer=model_name)

# Load DFA Press Release dataset
df = pd.read_csv('dfa_pr_v5_cleaned.csv', nrows=500)

# Group into 6 sentences-long parts
partitions = df['article'].values.tolist()

st.title('DFA Press Releases - Question Answer Bot')

# Type in HP-related query here
query = st.text_area("Type in your question below:")

if st.button('Search for the answer'):
    # Perform sentence embedding on query and sentence groups
    model_embed_name = 'sentence-transformers/msmarco-distilbert-dot-v5'
    
    model_embed = SentenceTransformer(model_embed_name)
    doc_emb = model_embed.encode(partitions)
    query_emb = model_embed.encode(query)
    
    #Compute dot score between query and all document embeddings
    scores = util.dot_score(query_emb, doc_emb)[0].cpu().tolist()
    
    #Combine docs & scores
    doc_score_pairs = list(zip(partitions, scores))
    
    #Sort by decreasing score and get only 3 most similar groups
    doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1],
                         reverse=True)[:1]
    
    # Join these similar groups to form the context 
    context = "".join(x[0] for x in doc_score_pairs)

    # Perform the querying
    QA_input = {'question': query, 'context': context}
    res = pipe(QA_input)
    
    confidence = res.get('score')
    if confidence > 0.8:
        st.write(res.get('answer'))
    else:
        out = "I am not sure."
        st.write(out)
    #out = res.get('answer')