app.py

import gradio as gr
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline

def load_model(model_name):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    return pipeline('text-classification', model=model, tokenizer=tokenizer, truncation=True, max_length=512, top_k=4)

classifier = load_model("ngocminhta/authscan-baseline")
classifier2 = load_model("ngocminhta/authscan-baseline-machine")

TEXT_CLASS_MAPPING_MACHINE = {
    'LABEL_0': 'Gemini 1.5 Pro',
    'LABEL_1': 'Gemini 2.0 Experimental',
    'LABEL_2': 'GPT-4o Mini',
    'LABEL_3': 'Llama 3.1 8B'
}

TEXT_CLASS_MAPPING = {
    'LABEL_0': 'Human-Written',
    'LABEL_1': 'Machine-Generated'
}

def update_language(language):
    if language == 'Java' or language == 'Python':
        return gr.update(language='python')
    elif language == 'C':
        return gr.update(language='c')
    elif language == 'C++':
        return gr.update(language='cpp')
    return gr.update(language='python')

def process_result_detection_tab(text, language): 
    result = classifier(f"Language: {language}\n\n{text}")[0]
    result_machine = classifier2(f"Language: {language}\n\n{text}")[0]

    labels = [TEXT_CLASS_MAPPING[x['label']] for x in result]
    labels_machine = [TEXT_CLASS_MAPPING_MACHINE[x['label']] for x in result_machine]

    scores = list(np.array([x['score'] for x in result]))
    scores_machine = list(np.array([x['score'] for x in result_machine]))
    final_results = dict(zip(labels, scores))
    if max(final_results, key=final_results.get) == 'Machine-Generated':
        final_results_machine = dict(zip(labels_machine, scores_machine))
    else:
        final_results_machine = None
    return final_results, final_results_machine

def clear_detection_tab():
    return "", gr.update(interactive=False)

css = """
body.dark-mode {
    background-color: #2c2c2c;
    color: white;
}
body.light-mode {
    background-color: white;
    color: black;
}
#dark-mode-toggle {
    position: fixed;
    top: 10px;
    right: 10px;
    z-index: 1000;
    padding: 10px 15px;
    background-color: #007bff;
    color: white;
    border: none;
    border-radius: 5px;
    cursor: pointer;
    font-size: 14px;
}
#dark-mode-toggle:hover {
    background-color: #0056b3;
}
"""

js = """
function toggleDarkMode() {
    const body = document.body;
    if (body.classList.contains('dark-mode')) {
        body.classList.remove('dark-mode');
        body.classList.add('light-mode');
    } else {
        body.classList.remove('light-mode');
        body.classList.add('dark-mode');
    }
}
"""

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""<h1><center>AuthScan</center></h1>""")
    gr.HTML("<script>{}</script>".format(js))
    gr.HTML(
        '<button id="dark-mode-toggle" onclick="toggleDarkMode()">Switch</button>'
    )
    with gr.Tab("Code Detection"):
        with gr.Row():
            language = gr.Dropdown(
                choices=["C", "C++", "Java", "Python"],
                label="Select Programming Language",
                value="C"
            )
    
        with gr.Row():
            input_text = gr.Code(
              label="Enter code here",
              language="python",
              elem_id="code_input",
            )
    
        with gr.Row():
            check_button = gr.Button("Check Origin", variant="primary")
            clear_button = gr.Button("Clear", variant="stop")
    
        out = gr.Label(label='Result')
        out_machine = gr.Label(label='Detailed Information')
        # When language is changed, update the code component's language
        language.change(update_language, inputs=language, outputs=input_text)
    
        check_button.click(process_result_detection_tab, inputs=[input_text, language], outputs=[out, out_machine])
        # out_machine.change(lambda x: gr.update(visible=True) if out_machine else gr.update(visible=False), inputs=out_machine, outputs=out_machine)
        clear_button.click(clear_detection_tab, inputs=[], outputs=[input_text, check_button])
        
    with gr.Tab("Text Detection"):
        gr.Markdown("""Under development!""")
        
demo.launch(share=True)