import streamlit as st
from backend_model import load_model_and_tokenizer, infer_single_sample


java_model_architecture = 'microsoft/graphcodebert-base'
java_model_path = 'models/java_classifier.pth'

python_model_architecture = 'microsoft/graphcodebert-base'
python_model_path = 'models/python_classifier.pth'

@st.cache_resource
def load_model(arch, path):
    return  load_model_and_tokenizer(arch, path)



st.title('LLM Sniffer')

# form
with st.form(key='my_form'):


    # select language - java or python
    language = st.selectbox(
        label="Select Language",
        options=["Java", "Python"],
        key="language"
    )

    
    # text area

    code = st.text_area(label="", value="", label_visibility="hidden", height=300, placeholder="Paste your code here", key="code")

    # submit button
    submit_button = st.form_submit_button(label='Submit')

    if submit_button:
        if code:

            if language == "Java":
                model, tokenizer = load_model(java_model_architecture, java_model_path)
            else:
                model, tokenizer = load_model(python_model_architecture, python_model_path)
            
            result = infer_single_sample(
                code_text=code, 
                model=model, 
                tokenizer=tokenizer,
                language=language
            )
            st.write(result)