Update app.py

app.py

@@ -1,15 +1,69 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
 import gradio as gr
-import subprocess
 
-def merge_model():
-    subprocess.run(["python", "merge_and_save_model.py"], check=True)
-    return "Model merged and saved successfully!"
+# Use GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-with gr.Blocks() as demo:
-    gr.Markdown("<h1>🧠 Phi-2 QLoRA Model Merger</h1>")
-    merge_button = gr.Button("Merge Model")
-    output = gr.Textbox(label="Merge Status")
+# Base model and adapter paths
+base_model_name = "microsoft/phi-2"  # Pull from HF Hub directly
+adapter_path = "Shriti09/Microsoft-Phi-QLora"  # Update with your Hugging Face repo path
 
-    merge_button.click(merge_model, [], output)
+print("🔧 Loading base model...")
+# Load the base model
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
+)
 
+print("🔧 Loading LoRA adapter...")
+# Load the LoRA adapter
+adapter_model = PeftModel.from_pretrained(base_model, adapter_path)
+
+print("🔗 Merging adapter into base model...")
+# Merge adapter into the base model
+merged_model = adapter_model.merge_and_unload()
+merged_model.eval()
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+print("✅ Model ready for inference!")
+
+# Text generation function
+def generate_text(prompt):
+    # Tokenize the input
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        outputs = merged_model.generate(
+            **inputs,
+            max_new_tokens=150,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+    # Decode and return the generated response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
+
+# Gradio UI
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("<h1>🧠 Phi-2 QLoRA Text Generator</h1>")
+    
+    # Textbox for user input
+    prompt = gr.Textbox(label="Enter your prompt:", lines=2)
+    
+    # Output textbox for generated text
+    output = gr.Textbox(label="Generated text:", lines=5)
+    
+    # Button to trigger text generation
+    generate_button = gr.Button("Generate Text")
+
+    # Set the button action to generate text
+    generate_button.click(generate_text, inputs=prompt, outputs=output)
+
+# Launch the app
 demo.launch(share=True)