import gradio as gr
import torch
import spaces
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("ping98k/typhoon-7b-rag-instruct-th")
model = AutoModelForCausalLM.from_pretrained("ping98k/typhoon-7b-rag-instruct-th", device_map={"": 0})

'''@spaces.GPU(duration=120)
def response(instruction, history, inputText):
    inp = f"""### Instruction:
{instruction}

### Input:
=======START OF DOCUMENT=======
{inputText}
=======END OF DOCUMENT=======

### Response:"""
    print(f"Ask: {instruction}")
    
    input_ids = tokenizer(inp, return_tensors='pt').to("cuda")
    beam_output = model.generate(**input_ids, max_new_tokens=300)
    outputText = tokenizer.decode(beam_output[0], skip_special_token=True)
    outputText = outputText.replace(inp,"").replace("<s>","").replace("</s>","")

    print(f"Ans: {outputText}")
    
    return outputText

defaultInput = """วัดอรุณราชวรารามราชวรมหาวิหาร หรือที่เรียกสั้น ๆ ว่า วัดอรุณ เป็นวัดโบราณ สร้างในสมัยอยุธยา ที่ชื่อวัดแจ้ง เพราะ พระเจ้าตากฯ ทำศึกเสร็จ แล้วยกทัพกลับมาเป็นเวลาเช้าพอดี ว่ากันว่าเดิมเรียกว่า วัดมะกอก และกลายเป็นวัดมะกอกนอกในเวลาต่อมา เพราะได้มีการสร้างวัดขึ้นอีกวัดหนึ่งในตำบลเดียวกัน แต่อยู่ในคลองบางกอกใหญ่ ชาวบ้านเรียกวัดที่สร้างใหม่ว่า วัดมะกอกใน (วัดนวลนรดิศ) แล้วจึงเรียกวัดมะกอกซึ่งอยู่ปากคลองบางกอกใหญ่ว่า วัดมะกอกนอก ส่วนเหตุที่มีการเปลี่ยนชื่อเป็นวัดแจ้งนั้น เชื่อกันว่า เมื่อสมเด็จพระเจ้ากรุงธนบุรีทรงตั้งราชธานีที่กรุงธนบุรีใน พ.ศ. 2310 ได้เสด็จมาถึงหน้าวัดนี้ตอนรุ่งแจ้ง จึงพระราชทานชื่อใหม่ว่าวัดแจ้ง
เมื่อสมเด็จพระเจ้ากรุงธนบุรีโปรดเกล้าฯ ให้สร้างพระราชวังที่ประทับนั้น ทรงเอาป้อมวิชัยประสิทธิ์ข้างฝั่งตะวันตกเป็นที่ตั้งตัวพระราชวัง แล้วขยายเขตพระราชฐานจนวัดแจ้งเป็นวัดภายในพระราชวัง เช่นเดียวกับวัดพระศรีสรรเพชญ์สมัยอยุธยา และเป็นที่ประดิษฐานพระพุทธมหามณีรัตนปฏิมากรที่อัญเชิญมาจากเวียงจันทน์ใน พ.ศ. 2321 ก่อนที่จะย้ายมาประดิษฐานที่วัดพระศรีรัตนศาสดารามในปี พ.ศ. 2327
ในสมัยรัตนโกสินทร์ รัชสมัยพระบาทสมเด็จพระพุทธยอดฟ้าจุฬาโลกมหาราช สมเด็จพระเจ้าลูกยาเธอ เจ้าฟ้ากรมหลวงอิศรสุนทร ได้เสด็จมาประทับที่พระราชวังเดิม และได้ทรงปฏิสังขรณ์วัดแจ้งใหม่ทั้งวัด แต่ยังไม่ทันสำเร็จก็สิ้นรัชกาลที่ 1 สมเด็จฯ เจ้าฟ้ากรมหลวงอิศรสุนทรได้เสด็จขึ้นครองราชสมบัติเป็นพระบาทสมเด็จพระพุทธเลิศหล้านภาลัย พระองค์ได้ทรงบูรณปฏิสังขรณ์วัดแจ้งต่อมา และพระราชทานนามใหม่ว่า “วัดอรุณราชธาราม” ต่อมามีพระราชดำริที่จะเสริมสร้างพระปรางค์หน้าวัดให้สูงขึ้น แต่สิ้นรัชกาลเสียก่อน จนถึงรัชสมัยพระบาทสมเด็จพระนั่งเกล้าเจ้าอยู่หัว โปรดเกล้าฯ ให้เสริมพระปรางค์ขึ้นและให้ยืมมงกุฎที่หล่อสำหรับพระพุทธรูปทรงเครื่องที่จะเป็นพระประธานวัดนางนองมาติดต่อบนยอดนภศูล ในรัชสมัยพระบาทสมเด็จพระจอมเกล้าเจ้าอยู่หัวโปรดเกล้าฯ ให้บูรณปฏิสังขรณ์วัดอรุณราชธารามหลายรายการ และให้อัญเชิญพระบรมอัฐิของพระบาทสมเด็จพระพุทธเลิศหล้านภาลัยมาบรรจุไว้ที่พระพุทธอาสน์ของพระประธานในพระอุโบสถด้วย เมื่อการปฏิสังขรณ์เสร็จสิ้นลง พระราชทานนามวัดใหม่ว่า วัดอรุณราชวราราม"""


gr.ChatInterface(
    response,
    additional_inputs=[
        gr.Textbox(defaultInput, label="Input Text"),
    ],
).launch()'''

@spaces.GPU(duration=120)
def response(inp):
    input_ids = tokenizer(inp, return_tensors='pt').to("cuda")
    beam_output = model.generate(**input_ids, max_new_tokens=300)
    outputText = tokenizer.decode(beam_output[0], skip_special_token=True)
    outputText = outputText.replace(inp,"").replace("<s>","").replace("</s>","")

    print(f"Ans: {outputText}")
    
    return outputText

gradio_interface = gr.Interface(
  fn=response,
  inputs = "text",
  outputs = "text"
)
gradio_interface.launch()