from unsloth import FastLanguageModel
import torch
import gradio as gr

max_seq_length = 2048 # Choose any! We auto support RoPE Scaling internally!
dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
load_in_4bit = True # Use 4bit quantization to reduce memory usage. Can be False.
alpaca_prompt = """Berikut adalah instruksi yang deskripsikan tugas dan sepasang input dan konteksnya. Tulis response sesuai dengan permintaan.
### Instruction:
{}
### Input:
{}
### Response:
{}"""

if True:
    from unsloth import FastLanguageModel
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name = "abdfajar707/llama3_8B_lora_model_rkp_pn2025_v3", # YOUR MODEL YOU USED FOR TRAINING
        max_seq_length = max_seq_length,
        dtype = dtype,
        load_in_4bit = load_in_4bit,
    )
    FastLanguageModel.for_inference(model) # Enable native 2x faster inference



# Fungsi untuk menghasilkan respons
def generate_response(prompt, max_length=1000):
    inputs = tokenizer(
[
    alpaca_prompt.format(
        prompt, # instruction
        "", # input
        "", # output - leave this blank for generation!
    )
], return_tensors = "pt").to("cuda")
    outputs = model.generate(**inputs, max_length=max_length, pad_token_id=tokenizer.eos_token_id)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

history = []
def wrapper_chat_history(chat_history, history):
    chat_history = history[1:]
    return chat_history

def converse(message, chat_history):
    response = generate_response(message)
    print(response)
    user_msg = {"role": "user", "content": message}
    history.append(user_msg)
    ai_msg = {"role": "assistant", "content": response}
    history.append(ai_msg)
    return history[-1]["content"]


DESCRIPTION = '''
<div style="padding: 5px; text-align: left; display: flex; flex-direction: column; align-items: left;">
   <img src="https://sdgs.bappenas.go.id/repository/assets/bappenas_logo_square.png" style="width: 40%; max-width: 200px; height: auto; opacity: 0.55;  "> 
   <h2 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">AI-Interlinked System/Bappenas GPT</h2>
</div>
'''

LICENSE = """
<p/>
---
Dibangun dari Meta Llama 3
"""

PLACEHOLDER = """
<div style="padding: 100px; text-align: center; display: flex; flex-direction: column; align-items: center;">
   <img src="https://cdn3.iconfinder.com/data/icons/human-resources-flat-3/48/150-4096.png" style="width: 1000; max-width: 200px; height: auto; opacity: 0.55;  "> 
   <h2 style="font-size: 20px; margin-bottom: 2px; opacity: 0.55;">Asisten Virtual Perencana</h2>
    <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Silakan mulai tanya...</p>
</div>
"""


css = """
h1 {
  text-align: center;
  display: block;
}
#duplicate-button {
  margin: auto;
  color: white;
  background: #1565c0;
  border-radius: 100vh;
}
"""
chatbot=gr.Chatbot(height=600, placeholder=PLACEHOLDER, label='Interlinked Sytem ChatInterface')
    
with gr.Blocks(css=css) as interface:
  chatbot=chatbot,
  with gr.Row():
    with gr.Column(scale=1):
      gr.HTML('<img src="https://datahub.data.go.id/data/static/Kementerian%20PPN%20Bappenas%20Tanpa%20Teks.png" width="100px" alt="Image" style="max-width: 100%;">')
  with gr.Row():
    with gr.Column(scale=1, elem_id='col'):
      gr.ChatInterface(fn=converse, title=("""
<center> 
<h1>KemenPPN/Bappenas</h1>
<b>AI-Interlinked System/Bappenas GPT<b>
</center>
"""
))

interface.launch()