import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
import gradio as gr
import re
# # ตั้งค่า paths และ quantization
model_path = "EXt1/KMUTT-CPE35-Typhoon-7B-news-reasoning"
base_model_name = "scb10x/typhoon-7b"

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True
)

base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    quantization_config=quantization_config,
)
model = PeftModel.from_pretrained(base_model,model_path)
tokenizer = AutoTokenizer.from_pretrained(base_model_name)

def generate_reasoning(title_text, label):
    prompt = f"""[Instruction]: วิเคราะห์และอธิบายว่าข่าวนี้เป็นข่าวจริงหรือเท็จ พร้อมเหตุผลประกอบแบบเป็นขั้นตอน 
    [News Title]: {title_text}
    [Label]: {label}

    [Reasoning]:"""

    # เข้ารหัส prompt
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    # สร้างคำตอบ
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,        
            temperature=0.65,          
            top_p=0.9,
            do_sample=True,
            repetition_penalty=1.08,     
            no_repeat_ngram_size=2,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    # ถอดรหัสเป็นข้อความ
    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # ===== ตัดเฉพาะ Reasoning =====
    if "[Reasoning]:" in full_response:
        reasoning_part = full_response.split("[Reasoning]:")[-1].strip()
    else:
        reasoning_part = full_response.strip()

    # ===== ตัดตรงที่เริ่ม hallucinate =====
    # ใช้ regular expression เพื่อตัดคำที่ไม่ต้องการ
    reasoning_part = re.split(r"News Title:|Instruction:|Label:", reasoning_part)[0].strip()

    return reasoning_part

iface = gr.Interface(
    fn=generate_reasoning,
    inputs=[
        gr.Textbox(label="News Title", placeholder="ใส่ชื่อข่าวที่ต้องการวิเคราะห์", lines=2),
        gr.Textbox(label="Label", placeholder="ใส่ label เช่น 'ข่าวจริง' หรือ 'ข่าวปลอม'", lines=1)
    ],
    outputs=gr.Textbox(label="Reasoning"),
    live=True,
    title="Thai News Fact-Checking",  # ชื่อแอปพลิเคชัน
    description="กรอกข่าวและ label เพื่อให้ระบบวิเคราะห์ว่าเป็นข่าวจริงหรือเท็จ พร้อมเหตุผลประกอบ"
)

# รันแอปพลิเคชัน
iface.launch()