Update app.py

app.py

@@ -1,38 +1,44 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
 import time
 
 # ----------------------------
-# 🚀 Load Model (CPU)
+# 🔹 Load base model + LoRA weights
 # ----------------------------
-model_name = "umarfarzan/my-finetuned-model2-lora"
+BASE_MODEL = "unsloth/qwen2.5-7b"  # Original base model
+LORA_WEIGHTS = "umarfarzan/my-finetuned-model2-lora"
 
-# LoRA + 4-bit config
-bnb_config = BitsAndBytesConfig(load_in_4bit=True)
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-print("Loading model...")
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    device_map="cpu",
-    quantization_config=bnb_config
-)
-model.eval()
-print("✅ Model loaded successfully!")
+@torch.inference_mode()
+def load_model():
+    print("Loading base model...")
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+    model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        device_map={"": device},
+        torch_dtype=torch.float32
+    )
+    print("Applying LoRA weights...")
+    model = PeftModel.from_pretrained(model, LORA_WEIGHTS, device_map={"": device})
+    model.eval()
+    print("✅ Model loaded successfully!")
+    return model, tokenizer
+
+model, tokenizer = load_model()
 
 # ----------------------------
-# 💡 Generate Training Program
+# 🔹 Generation function
 # ----------------------------
 def generate_training_program(instruction, max_tokens=500, temperature=0.7, top_p=0.9):
-    prompt_text = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
-
+    prompt_text = f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.
 ### Instruction:
 {instruction}
-### Input:
 ### Response:
 """
-    inputs = tokenizer([prompt_text], return_tensors="pt").to("cpu")
+    inputs = tokenizer([prompt_text], return_tensors="pt").to(device)
 
     start_time = time.time()
     outputs = model.generate(
@@ -43,8 +49,7 @@ def generate_training_program(instruction, max_tokens=500, temperature=0.7, top_
         do_sample=True,
         use_cache=True
     )
-    generation_time = time.time() - start_time
-
+    gen_time = time.time() - start_time
     generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
 
     if "### Response:" in generated_text:
@@ -52,33 +57,35 @@ def generate_training_program(instruction, max_tokens=500, temperature=0.7, top_
     else:
         response = generated_text
 
-    return response, f"⏱️ Generated in {generation_time:.2f} seconds"
+    return response, f"⏱️ Generated in {gen_time:.2f} seconds"
 
 # ----------------------------
-# 🎨 Gradio UI
+# 🔹 Gradio UI
 # ----------------------------
 examples = [
-    ["Design a detailed 1-week training program titled 'The Leader's Blueprint for Strategic Problem-Solving' for mid-level to senior-level managers."],
-    ["Create a 3-day workshop on effective communication skills for remote teams."],
-    ["Develop a 5-day leadership bootcamp for new managers focusing on team management and conflict resolution."],
+    ["Design a 1-week training program 'The Leader's Blueprint' for mid-level managers and team leads."],
+    ["Create a 3-day workshop on effective communication for remote teams."],
+    ["Develop a 5-day leadership bootcamp for new managers."],
+    ["Design a half-day data-driven decision-making session for executives."],
+    ["Create a 2-week onboarding program for new software engineers."]
 ]
 
 with gr.Blocks() as demo:
-    gr.Markdown("## 🎯 AI Training Program Generator")
-    instruction_input = gr.Textbox(label="Training Program Description", lines=5)
-    max_tokens_slider = gr.Slider(500, 2000, value=500, step=100, label="Max Output Length")
+    gr.HTML("<h1 style='text-align:center'>🎯 AI Training Program Generator</h1>")
+    instruction_input = gr.Textbox(label="📝 Training Program Description", lines=5)
+    max_tokens_slider = gr.Slider(100, 8000, value=500, step=100, label="Max Output Length")
     temperature_slider = gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Creativity (Temperature)")
     top_p_slider = gr.Slider(0.5, 1.0, value=0.9, step=0.05, label="Diversity (Top-p)")
     generate_btn = gr.Button("🚀 Generate Training Program")
-    output_text = gr.Textbox(label="Generated Training Program", lines=25, show_copy_button=True)
-    generation_info = gr.Textbox(label="Generation Info", interactive=False, show_label=False)
+    output_text = gr.Textbox(label="📋 Generated Training Program", lines=25, show_copy_button=True)
+    generation_info = gr.Textbox(label="ℹ️ Generation Info", interactive=False, show_label=False)
 
     generate_btn.click(
-        fn=generate_training_program,
+        generate_training_program,
         inputs=[instruction_input, max_tokens_slider, temperature_slider, top_p_slider],
         outputs=[output_text, generation_info]
     )
 
-    gr.Examples(examples=examples, inputs=[instruction_input])
+    gr.Examples(examples=examples, inputs=instruction_input)
 
-demo.launch(share=True)
+demo.launch(server_name="0.0.0.0", server_port=7860)