import gradio as gr
import re
import torch

def predict_house_price(description, location, area):
    FastLanguageModel.for_inference(model)  # Enable native 2x faster inference
    messages = [  # Change below!
        {
            "role": "user",
            "content": f"Mô tả: {description}\nĐịa chỉ: {location}\nDiện tích: {area}",
        },
    ]
    input_ids = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
    ).to("cuda")

    # from transformers import TextStreamer
    # text_streamer = TextStreamer(tokenizer, skip_prompt = True)
    # _ = model.generate(input_ids, streamer = text_streamer, max_new_tokens = 128, pad_token_id = tokenizer.eos_token_id)
    # return extract_price(_)
    output = model.generate(
        input_ids, max_new_tokens=256, pad_token_id=tokenizer.eos_token_id
    )

    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
    return f"{extract_price(decoded)} tỷ"


def extract_price(decoded):
    # match number like 6.950.000.000, or 6950000000, or 6.95B, etc.
    price_pattern = r"[\d\.\,]+"
    matches = re.findall(price_pattern, decoded)
    last_item = matches[-1]
    return last_item


from unsloth import FastLanguageModel

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.

# 4bit pre quantized models we support for 4x faster downloading + no OOMs.
fourbit_models = [
    "unsloth/mistral-7b-v0.3-bnb-4bit",  # New Mistral v3 2x faster!
    "unsloth/mistral-7b-instruct-v0.3-bnb-4bit",
    "unsloth/llama-3-8b-bnb-4bit",  # Llama-3 15 trillion tokens model 2x faster!
    "unsloth/llama-3-8b-Instruct-bnb-4bit",
    "unsloth/llama-3-70b-bnb-4bit",
    "unsloth/Phi-3-mini-4k-instruct",  # Phi-3 2x faster!
    "unsloth/Phi-3-medium-4k-instruct",
    "unsloth/mistral-7b-bnb-4bit",
    "unsloth/gemma-7b-bnb-4bit",  # Gemma 2.2x faster!
]  # More models at https://huggingface.co/unsloth

model, tokenizer = FastLanguageModel.from_pretrained(
    # model_name = "unsloth/llama-3-8b-bnb-4bit",
    model_name="QuangHoDev/lora_model",
    max_seq_length=max_seq_length,
    dtype=dtype,
    load_in_4bit=load_in_4bit,
    # token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
)

model = FastLanguageModel.get_peft_model(
    model,
    r=16,  # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],
    lora_alpha=16,
    lora_dropout=0,  # Supports any, but = 0 is optimized
    bias="none",  # Supports any, but = "none" is optimized
    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
    use_gradient_checkpointing="unsloth",  # True or "unsloth" for very long context
    random_state=3407,
    use_rslora=False,  # We support rank stabilized LoRA
    loftq_config=None,  # And LoftQ
)

from datasets import load_dataset

dataset = load_dataset("QuangHoDev/house-prices-info", split="train")

from unsloth import to_sharegpt

dataset = to_sharegpt(
    dataset,
    merged_prompt="Tiêu đề: {title}\nMô tả: {description}\nĐịa chỉ: {location}\nDiện tích: {area}",
    output_column_name="price",
    conversation_extension=3,  # Select more to handle longer conversations
)

from unsloth import standardize_sharegpt

dataset = standardize_sharegpt(dataset)

chat_template = """Dưới đây là thông tin về các bất động sản và giá của chúng.
Hãy đoán giá bất động sản theo mô tả.

>>> Mô tả bất động sản:
{INPUT}

>>> Giá là:
{OUTPUT}"""

from unsloth import apply_chat_template

dataset = apply_chat_template(
    dataset,
    tokenizer=tokenizer,
    chat_template=chat_template,
    # default_system_message = "You are a helpful assistant", << [OPTIONAL]
)

with gr.Blocks() as demo:

    gr.Markdown("""
    ## 🏠 Dự đoán giá nhà
    Nhập mô tả, địa chỉ và diện tích căn nhà để mô hình dự đoán giá.
    """)

    description = gr.Textbox(label="Mô tả nhà", lines=4)
    location = gr.Textbox(label="Địa chỉ", lines=2)
    area = gr.Number(label="Diện tích (m²)")
    output = gr.Textbox(label="Giá nhà")
    greet_btn = gr.Button("Đoán giá", variant='primary')
    greet_btn.click(fn=predict_house_price, inputs=[description, location, area], outputs=output, api_name="greet")

    gr.Examples(
        examples=[
            [
                "Nhanh tay sở hữu căn nhà riêng (căn góc) 3 tầng tại Phường 10, Gò Vấp...",
                "261/, Đường Quang Trung, Phường 10, Quận Gò Vấp",
                40
            ],
        ],
        inputs=[description, location, area],
        label="Ví dụ mẫu"
    )


    demo.launch()