from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
import gradio as gr
import torch

scam_model_name = "Tlezz324/thai-scam-detector-v1.69"

tokenizer = AutoTokenizer.from_pretrained(scam_model_name)
scam_model = AutoModelForSequenceClassification.from_pretrained(scam_model_name)

asr = pipeline("automatic-speech-recognition", model="airesearch/wav2vec2-large-xlsr-53-th")

def transcribe_and_predict(audio):
    if audio is None:
        return "No audio input detected", ""
    
    text = asr(audio)["text"]
    
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        logits = scam_model(**inputs).logits
        pred = torch.argmax(logits, dim=1).item()
    label = "Scam" if pred == 1 else "Not Scam"
    
    return text, label

iface = gr.Interface(
    fn=transcribe_and_predict,
    inputs=gr.Audio(type="filepath"),
    outputs=["text", "text"],
    title="Thai Scam Detector with Speech-to-Text",
    description="อัปโหลดไฟล์เสียงเพื่อแปลงเป็นข้อความและตรวจสอบว่าหลอกลวงหรือไม่"
)

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