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 nếu có
    if system_message:
        messages.append({"role": "system", "content": system_message})
    
    # 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
    )
    
    # 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()