Update app.py

app.py

@@ -3,7 +3,6 @@ import torch
 import gradio as gr
 from huggingface_hub import login
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel
 
 # === STEP 1: Authenticate with Hugging Face ===
 # Make sure you set your HF token as an environment variable or paste it here temporarily
@@ -13,32 +12,41 @@ login(token="hf_" + hf_token)
 
 # === STEP 2: Load base and adapter models ===
 base_model = "meta-llama/Llama-2-7b-chat-hf"
-adapter_model = "olacode55/zimble-llama2"
+adapter_model = "zimble-llama2-finetunedhybride"
 
-tokenizer = AutoTokenizer.from_pretrained(base_model)
+tokenizer = AutoTokenizer.from_pretrained(merged_model_repo, use_auth_token=hf_token)
 
-offload_folder = "./offload"  # must exist or be creatable
-os.makedirs(offload_folder, exist_ok=True)
+# Enable memory-efficient loading if needed
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# --- Load model with 8-bit quantization and CPU offload ---
-base = AutoModelForCausalLM.from_pretrained(
-    base_model,
-    load_in_8bit=True,
+model = AutoModelForCausalLM.from_pretrained(
+    merged_model_repo,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
     device_map="auto",
-    offload_folder=offload_folder,
-    llm_int8_enable_fp32_cpu_offload=True,
-    use_auth_token="hf_" +hf_token
+    low_cpu_mem_usage=True,
+    use_auth_token=hf_token
 )
 
-
-model = PeftModel.from_pretrained(base, adapter_model)
-
 # === STEP 3: Define generation function ===
 def generate(prompt):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=200, temperature=0.7)
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=250,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+        )
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # === STEP 4: Launch Gradio app ===
-demo = gr.Interface(fn=generate, inputs="text", outputs="text", title="Zimble LLaMA 2 Fine-Tuned")
+demo = gr.Interface(
+    fn=generate,
+    inputs=gr.Textbox(label="Enter your prompt", lines=4, placeholder="Type something..."),
+    outputs=gr.Textbox(label="Model output"),
+    title="🦙 Zimble LLaMA 2 (Merged)",
+    description="Fine-tuned and merged version of LLaMA 2 running on Hugging Face Space"
+)
+
 demo.launch()