import time
from io import StringIO
import streamlit as st
import joblib
from transformers import pipeline
from lmqg import TransformersQG

#util func starts here

def break_paragraph_into_parts(paragraph, max_length):
    sentences = paragraph.split(". ")  # Assuming that the sentences end with a period and a space
    temp_parts = []
    part = ''
    for sentence in sentences:
        if len(part) + len(sentence) <= max_length:
            part += sentence + ". "  # Add the sentence and the period back
        else:
            temp_parts.append(part)
            part = sentence + ". "  # Start a new part with the remaining sentence
    temp_parts.append(part)  # Add the last part

    parts = [part.strip() for part in temp_parts]  # Remove any leading/trailing spaces

    return parts


def util(NumQues,Input):

    context = break_paragraph_into_parts(Input,512)
    context = context[:NumQues]

    #Question generation
    generatedQuestions = []
    Question_Generator=joblib.load("Qgenerator.sav")

    for part in context:
        question = Question_Generator.generate_q(list_context=part, list_answer="")
        generatedQuestions.append(question)


    #Answer Generation
    load_pipeline=joblib.load('Agenerator.sav')

    generatedAnswers=[]
    for Q in generatedQuestions:
        print(Q,'\n')
        gen_answer=load_pipeline(question=Q, context=Input)
        generatedAnswers.append(gen_answer['answer'])

    for i in range(len(generatedAnswers)):
        code = f'Ques: "{generatedQuestions[i]}"\nAns: "{generatedAnswers[i]}"'
        st.code(code, language='python')

#util func ends here






st.title('Question Answer Pair Generation from Documents')

tab1, tab2 = st.tabs(["Enter Text", "Choose Document"])

flag='None'

with st.sidebar:
    st.image('Pic.png')
    st.title("Final Year Project")
    st.divider()
    code = '''Team Members CSE(20-37):
    \nPrateek Niket BT20CSE211 \nSmriti Singh BT20CSE156 \nAmbuj Raj BT20CSE054 \nSrishti Pandey BT20CSE068'''
    st.code(code, language='JAVA')
    code = '''Mentored By: \nDr. Amol Bhopale'''
    st.code(code, language='JAVA')


with tab1:
    txt = st.text_area(
        "Enter Text to Generate Question-Answer"
    )
    flag='text'

with tab2:
    uploaded_file = st.file_uploader("Choose a file", type=['txt'], accept_multiple_files=False)
    if uploaded_file is not None:
        # To convert to a string based IO:
        stringio = StringIO(uploaded_file.getvalue().decode("utf-8"))
        txt = stringio.read()
    flag='file'

NumQues = st.slider('No. of Questions to Generate: ', 1, 5, 1)


if st.button('Generate',type="primary"):
    with st.spinner('Question Answer pair Generation in Progress....'):
      util(NumQues,txt)
    st.success('Question Answer pair Generated Successfully!')