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


# Load tokenizer và model
print("Đang tải model...")
model_name = "cochi1706/codingassistant"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Xác định device cho pipeline (0 cho cuda, -1 cho cpu)
device = 0 if torch.cuda.is_available() else -1

# Set padding token nếu chưa có
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# Tạo pipeline để sinh text
text_generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=device,
    max_length=520,
    do_sample=True,
    temperature=0.7
)

model.eval()
print(f"Model đã sẵn sàng! Device: {device}")


def respond(
    message,
    history=None
):
    """
    Tạo phản hồi từ model coding assistant sử dụng pipeline
    """
    # Xử lý message - có thể là dict hoặc string
    if isinstance(message, dict):
        user_message = message.get("content", "")
    else:
        user_message = str(message) if message else ""
    
    # Nếu không có message, trả về thông báo lỗi
    if not user_message:
        return "Xin lỗi, tôi không nhận được câu hỏi của bạn."
    
    # Sử dụng pipeline để generate text
    generated = text_generator(user_message, num_return_sequences=1)
    
    # Lấy câu trả lời từ kết quả
    answer = generated[0]['generated_text']
    
    # Đảm bảo trả về string không rỗng
    if not answer or len(answer.strip()) == 0:
        answer = "Xin lỗi, tôi không thể tạo phản hồi."
    
    return answer


"""
Chatbot hỗ trợ lập trình sử dụng Qwen3 fine-tuned model
"""
chatbot = gr.ChatInterface(
    respond,
    title="🤖 Coding Assistant",
    description="Chatbot hỗ trợ lập trình",
    type="messages"
)

demo = chatbot


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