import streamlit as st
from transformers import pipeline
from ModelDriver import *
import numpy as np

# Add a title
st.title('GPT Detection Demo')
st.write("This is a demo for GPT detection. You can use this demo to test the model. The model is trained on two datasets: OpenGPT and CSAbstract. You can choose the model and dataset in the sidebar.")

# Add 4 options for 4 models
ModelOption = st.sidebar.selectbox(
    'Which Model do you want to use?',
    ('RobertaSentinel', 'RobertaClassifier'),
)

DatasetOption = st.sidebar.selectbox(
    'Which Dataset the model was trained on?',
    ('OpenGPT', 'CSAbstract'),
)


text = st.text_area('Enter text here (max 500 words)', '')

if st.button('Generate'):
    if ModelOption == 'RobertaSentinel':
        if DatasetOption == 'OpenGPT':
            result = RobertaSentinelOpenGPTInference(text)
            st.write("Model: RobertaSentinelOpenGPT")
        elif DatasetOption == 'CSAbstract':
            result = RobertaSentinelCSAbstractInference(text)
            st.write("Model: RobertaSentinelCSAbstract")

    elif ModelOption == 'RobertaClassifier':
        if DatasetOption == 'OpenGPT':
            result = RobertaClassifierOpenGPTInference(text)
            st.write("Model: RobertaClassifierOpenGPT")
        elif DatasetOption == 'CSAbstract':
            result = RobertaClassifierCSAbstractInference(text)
            st.write("Model: RobertaClassifierCSAbstract")

    Prediction = "Human Written" if not np.argmax(result) else "Machine Generated"

    st.write(f"Prediction: {Prediction} ")
    st.write(f"Probabilty:", max(result))