Update app.py

app.py

@@ -1,16 +1,16 @@
+
 # import gradio as gr
 # from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 # # --- Load Model ---
-# MODEL_PATH = "./tinyllama-jobskills-final_update_4"  # Model files are in the repo root
+# 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",
-#     low_cpu_mem_usage=True
-
+#     device_map="auto",   # Use GPU if available
+#     low_cpu_mem_usage=True,
 # )
 
 # pipe = pipeline(
@@ -20,41 +20,47 @@
 #     device_map="auto"
 # )
 
-# # --- Define Chat Function ---
+# # --- Chat Function ---
 # def chat_fn(message, history):
-#     history_text = ""
-#     for user, bot in history:
-#         history_text += f"User: {user}\nAssistant: {bot}\n"
-#     history_text += f"User: {message}\nAssistant:"
+#     # Use the same format as training data
+#     prompt = f"### Question:\n{message}\n\n### Answer:\n"
 
-#     # generate response
+#     # Generate response
 #     response = pipe(
-#         history_text,
-#         max_new_tokens=256,
-#         do_sample=True,
+#         prompt,
+#         max_new_tokens=32,   # allow longer output
+#         do_sample=False,
 #         temperature=0.7,
-#         top_p=0.9
+#         top_p=1.0
 #     )[0]["generated_text"]
 
-#     # extract assistant reply
-#     reply = response.split("Assistant:")[-1].strip()
-#     return reply
+#     # 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 Custom Model")
+#     gr.Markdown("## 🚀 Chat with My AI Skills Model")
 
-#     chatbot = gr.Chatbot()
-#     msg = gr.Textbox(label="Type your message")
+#     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((user_message, bot_message))
+#         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, None, chatbot, queue=False)
+#     clear.click(lambda: [], None, chatbot, queue=False)
 
 # # --- Launch ---
 # if __name__ == "__main__":
@@ -63,66 +69,80 @@
 
 
 
-
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
+import torch # Needed for torch.bfloat16
 
 # --- Load Model ---
-MODEL_PATH = "./tinyllama-jobskills-final_update_4"  # Path to your 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",   # Use GPU if available
+    device_map="auto",
     low_cpu_mem_usage=True,
-    load_in_4bit=True
+    quantization_config=bnb_config, # <--- Pass the BitsAndBytesConfig object here
 )
 
 pipe = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    device_map="auto"
+    device_map="auto" # Redundant if model is already on device, but harmless
 )
 
 # --- Chat Function ---
-def chat_fn(message, history):
-    # Use the same format as training data
+def chat_fn(message):
     prompt = f"### Question:\n{message}\n\n### Answer:\n"
 
-    # Generate response
     response = pipe(
         prompt,
-        max_new_tokens=32,   # allow longer output
+        max_new_tokens=32,
         do_sample=False,
-        temperature=0.7,
-        top_p=1.0
+        # 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"]
 
-    # 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...")
+    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):
-        bot_message = chat_fn(user_message, chat_history)
-        chat_history.append({"role": "user", "content": user_message})
-        chat_history.append({"role": "assistant", "content": bot_message})
+        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])
@@ -132,5 +152,3 @@ with gr.Blocks() as demo:
 if __name__ == "__main__":
     demo.launch()
 
-
-