Update app.py

app.py

@@ -1,47 +1,56 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 model_id = "google/gemma-4-E2B"
 
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+# Load model in LOW MEMORY MODE
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    device_map={"": "cpu"},
-    dtype=torch.float32,
+    device_map="auto",          # better memory distribution
+    torch_dtype=torch.float16,  # HUGE RAM saver vs float32
     low_cpu_mem_usage=True
 )
 
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device="cpu"
-)
+model.eval()
+
 
 def predict(message, history):
-    messages = [
-        {"role": "user", "content": message},
-    ]
-    
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    
-    outputs = pipe(
-        prompt,
-        max_new_tokens=512,
-        do_sample=True,
-        temperature=0.7,
-        top_k=50,
-        top_p=0.95,
+    messages = [{"role": "user", "content": message}]
+
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
     )
-    
-    return outputs[0]["generated_text"][len(prompt):]
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=256,   # lowered to reduce RAM spikes
+            do_sample=True,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.95,
+            use_cache=True
+        )
+
+    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    # return only new text
+    return decoded[len(prompt):]
+
 
 demo = gr.ChatInterface(
     fn=predict,
-    title="Gemma-4-E2B Chatbot",
-    description="Running on Free CPU Space - No Memory",
+    title="Gemma-4-E2B Chatbot (Optimized)",
+    description="Low RAM CPU-optimized version ⚡"
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()