import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch
from PIL import Image

# মডেল এবং টোকেনাইজার লোড করুন
model_name = "mistralai/Pixtral-12B-Base-2409"
model = GPT2LMHeadModel.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)

# BTroy চ্যাটবটের ফাংশন
def btroy_chat(input_text, temperature=1.0, top_p=1.0, frequency_penalty=0.0, input_image=None):
    # টেক্সট প্রসেসিং
    inputs = tokenizer.encode(input_text, return_tensors='pt')

    # ইমেজ প্রসেসিং যদি থাকে
    if input_image:
        image = Image.open(input_image)
        # এখানে আপনি ইমেজ প্রসেসিং কোড যোগ করতে পারেন (যেমন, ইমেজ থেকে ক্যাপশন জেনারেট করা)

    # মডেল আউটপুট
    outputs = model.generate(
        inputs, 
        max_length=200, 
        temperature=temperature, 
        top_p=top_p, 
        frequency_penalty=frequency_penalty
    )

    # আউটপুট টেক্সট ডিকোড করুন
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return response

# Gradio ইন্টারফেস তৈরি
iface = gr.Interface(
    fn=btroy_chat, 
    inputs=[
        gr.Textbox(label="আপনার প্রশ্ন দিন", placeholder="আপনার প্রশ্ন এখানে টাইপ করুন..."), 
        gr.Slider(minimum=0, maximum=2, default=1.0, label="Temperature"),
        gr.Slider(minimum=0, maximum=1, default=1.0, label="Top-p"),
        gr.Slider(minimum=0, maximum=2, default=0.0, label="Frequency Penalty"),
        gr.Image(label="আপনার ইমেজ দিন", type="pil", optional=True)
    ],
    outputs="text",
    title="BTroy AI চ্যাটবট",
    description="BTroy চ্যাটবট ব্যবহার করুন আপনার যেকোনো প্রশ্নের উত্তর পেতে। এটি মাল্টি-মোডাল ইনপুট সমর্থন করে (যেমন ইমেজ)।",
    theme="huggingface",
    allow_flagging="never"
)

# ইন্টারফেস রান করুন
iface.launch()