import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_PATH = "lzw1008/ConspEmoLLM-v2"
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForCausalLM.from_pretrained(MODEL_PATH)

def classify_conspiracy(text):
    prompt = f"""Human:
Task: Classify the text regarding conspiracy theories or misinformation.
Text: {text}
Assistant:
"""
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=128)
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result

demo = gr.Interface(
    fn=classify_conspiracy,
    inputs=gr.Textbox(
        lines=5,
        placeholder="Enter text to classify...",
        label="Input Text"
    ),
    outputs=gr.Textbox(
        lines=10,
        label="Classification Result"
    ),
    title="ConspEmoLLM-v2 API",
    description="Classifies input text for conspiracy theories using the ConspEmoLLM-v2 LLM."
)

if __name__ == "__main__":
    demo.launch()