app.py

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,
    do_sample=True,
)

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


def respond(
    message,
    history: list[dict[str, str]],
    system_message=None,
    max_tokens=None,
    temperature=None,
    top_p=None,
):
    """
    Tạo phản hồi từ model coding assistant sử dụng pipeline
    """
    # Đặt giá trị mặc định nếu không có
    if max_tokens is None:
        max_tokens = 250
    if temperature is None:
        temperature = 0.7
    if top_p is None:
        top_p = 0.95
    
    # 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 ""
    
    # Chuẩn bị messages từ history và message hiện tại
    messages = []
    
    # Thêm system message - mặc định bằng tiếng Việt nếu không có
    if system_message:
        messages.append({"role": "system", "content": system_message})
    else:
        # System message mặc định bằng tiếng Việt để model trả lời bằng tiếng Việt
        messages.append({
            "role": "system", 
            "content": "Bạn là một trợ lý lập trình hữu ích. Hãy trả lời bằng tiếng Việt. Cung cấp các giải pháp code rõ ràng, ngắn gọn và chính xác cùng với giải thích."
        })
    
    # Thêm history (đã có format đúng)
    messages.extend(history)
    
    # Thêm message hiện tại từ user
    if user_message:
        messages.append({"role": "user", "content": user_message})
    
    # Format messages với chat template
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    
    # Thêm instruction rõ ràng để model trả lời bằng tiếng Việt
    # (nếu model vẫn có xu hướng trả lời tiếng Anh)
    prompt += "\n\nLưu ý: Hãy trả lời bằng tiếng Việt."
    
    # Sử dụng pipeline để generate text - chỉ dùng max_new_tokens
    generated = text_generator(
        prompt,
        max_new_tokens=max_tokens,
        num_return_sequences=1,
        temperature=temperature,
        top_p=top_p,
        do_sample=True,
        truncation=True,
    )
    
    # Lấy câu trả lời từ kết quả
    answer = generated[0]['generated_text']
    
    # Loại bỏ prompt ban đầu để chỉ lấy phần response
    if prompt in answer:
        answer = answer.replace(prompt, "").strip()
    
    # Đả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()