import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# ১. আপনার মডেলের তথ্য
base_model_id = "unsloth/llama-3-8b-bnb-4bit"
adapter_model_id = "ZenJony/lora" # আপনার আপলোড করা আইডি

# ২. মডেল এবং টোকেনাইজার লোড করা
tokenizer = AutoTokenizer.from_pretrained(base_model_id)
model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    torch_dtype=torch.float16,
    device_map="auto", # এটি অটোমেটিক CPU বা GPU বেছে নেবে
)
# আপনার লরা অ্যাডাপ্টার যুক্ত করা
model = PeftModel.from_pretrained(model, adapter_model_id)

# ৩. আলপাকা প্রম্পট ফরম্যাট (ট্রেইনিং এর সময় যা ব্যবহার করেছিলেন)
alpaca_prompt = """তুমি একজন আধুনিক ও স্মার্ট এআই অ্যাসিস্ট্যান্ট। তোমার কাজ হলো মানুষের প্রশ্নের সঠিক ও সৃজনশীল উত্তর দেওয়া। উত্তরের গুরুত্ব বুঝে প্রাসঙ্গিক ইমোজি ব্যবহার করো এবং গুরুত্বপূর্ণ শব্দগুলো **বোল্ড** করো। যদি কোনো তথ্য না জানো, তবে বিনয়ের সাথে স্বীকার করো এবং বিকল্প পরামর্শ দাও।

### Instruction:
{}

### Input:
{}

### Response:
{}"""

def respond(
    message,
    history,
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    # ইনপুট এবং সিস্টেম মেসেজ একসাথে করা
    full_instruction = f"{system_message}\n\n{message}"
    
    # প্রম্পট তৈরি
    prompt = alpaca_prompt.format(full_instruction, "", "")
    inputs = tokenizer([prompt], return_tensors="pt").to(model.device)

    # উত্তর জেনারেট করা
    with torch.no_grad():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
            eos_token_id=tokenizer.eos_token_id,
        )
    
    # উত্তর ডিকোড করা
    full_response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    
    # শুধু Response অংশটুকু আলাদা করা
    if "### Response:" in full_response:
        response = full_response.split("### Response:")[1].strip()
    else:
        response = full_response
        
    return response

# ৪. চ্যাট ইন্টারফেস কাস্টমাইজেশন
chatbot = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="তুমি একজন আধুনিক ও স্মার্ট এআই অ্যাসিস্ট্যান্ট। তোমাকে তৈরি করেছেন ZenJony।", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
    ],
    title="ZenJony AI Assistant 🤖",
    description="আমার নিজের তৈরি ১০০০+ ডাটা দিয়ে ফাইন-টিউন করা বাংলা এআই মডেল।"
)

if __name__ == "__main__":
    chatbot.launch()