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

# # --- Load Model ---
# MODEL_PATH = "./tinyllama-jobskills-final_update_4"  # Path to your model

# tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
# model = AutoModelForCausalLM.from_pretrained(
#     MODEL_PATH,
#     trust_remote_code=True,
#     device_map="auto",   # Use GPU if available
#     low_cpu_mem_usage=True,
# )

# pipe = pipeline(
#     "text-generation",
#     model=model,
#     tokenizer=tokenizer,
#     device_map="auto"
# )

# # --- Chat Function ---
# def chat_fn(message, history):
#     # Use the same format as training data
#     prompt = f"### Question:\n{message}\n\n### Answer:\n"

#     # Generate response
#     response = pipe(
#         prompt,
#         max_new_tokens=32,   # allow longer output
#         do_sample=False,
#         temperature=0.7,
#         top_p=1.0
#     )[0]["generated_text"]

#     # Extract only the answer part
#     reply = response.split("### Answer:")[-1].strip()

#     # Format into bullet points
#     skills = [s.strip() for s in reply.replace(",", "\n").split("\n") if s.strip()]
#     formatted_reply = "\n".join([f"- {s}" for s in skills])

#     return formatted_reply



# # --- Gradio UI ---
# with gr.Blocks() as demo:
#     gr.Markdown("## 🚀 Chat with My AI Skills Model")

#     chatbot = gr.Chatbot(type="messages")
#     msg = gr.Textbox(label="Type your question here...")
#     clear = gr.Button("Clear")

#     def user_fn(user_message, chat_history):
#         bot_message = chat_fn(user_message, chat_history)
#         chat_history.append({"role": "user", "content": user_message})
#         chat_history.append({"role": "assistant", "content": bot_message})
#         return "", chat_history

#     msg.submit(user_fn, [msg, chatbot], [msg, chatbot])
#     clear.click(lambda: [], None, chatbot, queue=False)

# # --- Launch ---
# if __name__ == "__main__":
#     demo.launch()




import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
import torch # Needed for torch.bfloat16

# --- Load Model ---
MODEL_PATH = "./tinyllama-jobskills-final_update_4"  
# --- Define Quantization Configuration ---
# This is the new way to specify 4-bit or 8-bit loading
# For 4-bit:
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4", # Or "fp4"
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16, # Or torch.float16 if not using bfloat16
)

# For 8-bit (if preferred, though 4-bit is smaller and often good enough)
# bnb_config = BitsAndBytesConfig(
#     load_in_8bit=True
# )

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    trust_remote_code=True,
    device_map="auto",
    low_cpu_mem_usage=True,
    quantization_config=bnb_config, # <--- Pass the BitsAndBytesConfig object here
)

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device_map="auto" # Redundant if model is already on device, but harmless
)

# --- Chat Function ---
def chat_fn(message):
    prompt = f"### Question:\n{message}\n\n### Answer:\n"

    response = pipe(
        prompt,
        max_new_tokens=32,
        do_sample=False,
        # temperature and top_p are ignored when do_sample=False, so remove them:
        # temperature=0.7,
        # top_p=1.0,
        return_full_text=False # Get only the newly generated text
    )[0]["generated_text"]

    reply = response.split("### Answer:")[-1].strip()

    skills = [s.strip() for s in reply.replace(",", "\n").split("\n") if s.strip()]
    formatted_reply = "\n".join([f"- {s}" for s in skills])

    return formatted_reply


# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("## 🚀 Chat with My AI Skills Model")

    chatbot = gr.Chatbot(label="Chat History")
    msg = gr.Textbox(label="Type your question here...", placeholder="Ask about job skills...")
    clear = gr.Button("Clear")

    def user_fn(user_message, chat_history):
        chat_history = chat_history or []
        chat_history.append([user_message, None])

        bot_message = chat_fn(user_message)

        chat_history[-1][1] = bot_message
        return "", chat_history

    msg.submit(user_fn, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: [], None, chatbot, queue=False)

# --- Launch ---
if __name__ == "__main__":
    demo.launch()