Update app.py

added 4 bit quantization code
added examples

app.py

@@ -1,73 +1,102 @@
 import os
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from peft import PeftModel
 
-# Read your Hugging Face token from Space Secrets
+# Hugging Face Token from Space Secrets
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
-# Hugging Face model identifiers
+# Model IDs
 BASE_MODEL = "google/gemma-3-1b-it"
-LORA_ADAPTER = "markredito/gemma-pip-finetuned-v2"  # 🔁 
+LORA_ADAPTER = "markredito/gemma-pip-finetuned-v2"  # 🔁 Replace with your actual LoRA repo
 
-# Load base model with token (required for gated models)
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=HF_TOKEN)
-
-# Detect if GPU is available
+# Check device
 device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+# Quantization config for 4-bit (recommended on T4 GPU)
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+)
+
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=HF_TOKEN)
 
 model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
-    device_map="auto" if device == "cuda" else None,
-    torch_dtype=dtype,
-    token=HF_TOKEN
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+    quantization_config=bnb_config,
+    token=HF_TOKEN,
+    attn_implementation="eager"  # Required for Gemma3 + quant
 )
 
-model = PeftModel.from_pretrained(
-    model,
-    LORA_ADAPTER,
-    token=HF_TOKEN
-)
+model = PeftModel.from_pretrained(model, LORA_ADAPTER, token=HF_TOKEN)
 
 # Pad token fallback
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
     tokenizer.padding_side = "right"
 
-def generate_response(user_input):
-    prompt = (
+# Generation function
+def generate_response(prompt, temperature, top_p, top_k):
+    formatted = (
         "<start_of_turn>user\n"
-        f"{user_input.strip()}\n"
+        f"{prompt.strip()}\n"
         "<end_of_turn>\n"
         "<start_of_turn>model\n"
     )
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
+
+    inputs = tokenizer(formatted, return_tensors="pt").to(model.device)
+
     outputs = model.generate(
         **inputs,
         max_new_tokens=200,
         do_sample=True,
-        temperature=0.7,
-        top_p=0.9,
-        top_k=50,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
         pad_token_id=tokenizer.pad_token_id,
         eos_token_id=tokenizer.eos_token_id,
     )
 
-    # Decode and clean output
-    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
-    response = response.split("<end_of_turn>")[0].replace("model\n", "").strip()
-    
-    return response
+    decoded = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+    cleaned = decoded.split("<end_of_turn>")[0].replace("model\n", "").strip()
+    return cleaned
 
 # Gradio UI
-gr.Interface(
-    fn=generate_response,
-    inputs=gr.Textbox(label="Enter your prompt", placeholder="E.g. Describe a universe made of sound..."),
-    outputs=gr.Textbox(label="Model's response"),
-    title="Gemma LoRA: Abstract Thought Generator",
-    description="LoRA fine-tuned `gemma-3-1b-it` on poetic/philosophical prompts. Run your own abstract experiments.",
-    theme="soft"
-).launch()
+with gr.Blocks() as demo:
+    gr.Markdown("## ✨ Gemma LoRA Inference Demo")
+    gr.Markdown("Use your imagination or try one of the examples below to explore poetic and philosophical responses.")
+
+    examples = [
+        "Describe a world where clouds are solid and people walk on them",
+        "Contrast quantum realities phenomena from the perspective of a starship navigator, using a spiral into infinity.",
+        "Dream up futuristic phenomena from the perspective of a timeless oracle, using a fractal blooming in chaos.",
+    ]
+
+    with gr.Row():
+        with gr.Column():
+            prompt_input = gr.Textbox(label="Enter your prompt", lines=4, placeholder="Try something like: What if gravity took a day off?")
+
+            gr.Examples(
+                examples=examples,
+                inputs=prompt_input,
+                label="Example Prompts"
+            )
+
+            temperature = gr.Slider(0.1, 1.5, value=0.7, label="Temperature")
+            top_p = gr.Slider(0.1, 1.0, value=0.9, label="Top-p (nucleus sampling)")
+            top_k = gr.Slider(0, 100, step=1, value=50, label="Top-k")
+
+            submit = gr.Button("Generate")
+
+        with gr.Column():
+            output = gr.Textbox(label="Model Response", lines=10)
+
+    submit.click(fn=generate_response, inputs=[prompt_input, temperature, top_p, top_k], outputs=output)
+
+demo.launch()