import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# โหลดโมเดล Falcon-RW-1B
model_id = "tiiuae/falcon-rw-1b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

# ฟังก์ชันตอบกลับ
def ask_doctor(message, history):
    prompt = f"You are a friendly, professional doctor who gives clear, concise advice.\n\nPatient: {message}\nDoctor:"
    response = generator(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)[0]["generated_text"]
    answer = response.split("Doctor:")[-1].strip()
    answer += "\n\n⚠️ คำตอบนี้เพื่อการศึกษา ไม่แทนการวินิจฉัยจากแพทย์จริง"
    return answer

# สร้างหน้าตาแชตแบบ Chatbot
with gr.Blocks() as demo:
    gr.Markdown("""
        # 🩺 ปรึกษาหมอ AI
        ### หมอ AI พร้อมให้คำแนะนำเบื้องต้นด้านสุขภาพแบบรวดเร็ว ใช้โมเดล Falcon-RW-1B ฟรีจาก Hugging Face
        > ❗ **ข้อมูลนี้ใช้เพื่อการศึกษาเท่านั้น ไม่แทนการวินิจฉัยทางการแพทย์จริง**
    """)
    chatbot = gr.ChatInterface(
        fn=ask_doctor,
        title="แชตกับหมอ AI 🤖",
        chatbot=gr.Chatbot(height=400),
        textbox=gr.Textbox(placeholder="พิมพ์คำถาม เช่น 'ปวดหัวมาก ต้องทำยังไงดี?'", scale=7),
        theme="default"
    )

demo.launch()