Update app.py

app.py

@@ -1,44 +1,62 @@
-from transformers import pipeline
-import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
+import gradio as gr
+import os
 
-# Verify GPU
+# 1. GPU Verification
 if not torch.cuda.is_available():
-    raise RuntimeError("Enable GPU in Space settings")
+    raise RuntimeError("❌ GPU not detected! Enable GPU in Space settings.")
+print(f"✅ Using GPU: {torch.cuda.get_device_name(0)}")
+
+# 2. Model Configuration
+MODEL_NAME = "google/gemma-2b-it"  # Try "mistralai/Mistral-7B-v0.1" for more power
 
-# Load model (without bitsandbytes)
-model = pipeline(
+# 3. Load Model with GPU Optimization
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    device_map="auto",          # Auto-selects GPU
+    torch_dtype=torch.float16,  # Half-precision for memory
+    low_cpu_mem_usage=True      # Reduces CPU overhead
+)
+
+# 4. Create GPU-accelerated pipeline
+pipe = pipeline(
     "text-generation",
-    model="google/gemma-2b-it",
-    device=0,  # Force GPU
-    torch_dtype=torch.float16,
-    model_kwargs={
-        "low_cpu_mem_usage": True,
-        "trust_remote_code": True
-    }
+    model=model,
+    tokenizer=tokenizer,
+    device=0,                  # Force first GPU
+    torch_dtype=torch.float16
 )
 
-def solve_math(question):
-    prompt = f"Solve step by step:\nQ: {question}\nA:"
+# 5. Generation Function
+def generate_text(prompt):
     try:
-        result = model(
+        outputs = pipe(
             prompt,
-            max_new_tokens=100,
-            temperature=0.3,
-            do_sample=False
+            max_new_tokens=150,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
         )
-        return result[0]['generated_text'].split("A:")[-1].strip()
+        return outputs[0]['generated_text']
     except Exception as e:
-        return f"Error: {str(e)}"
-
-# Preload
-solve_math("2+2=")
+        return f"⚠️ Error: {str(e)}"
 
-# Simple UI
-with gr.Blocks() as demo:
-    gr.Markdown("## Math Solver")
-    question = gr.Textbox(label="Problem")
-    answer = gr.Textbox(label="Solution")
-    question.submit(solve_math, question, answer)
+# 6. Gradio Interface
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("## 🚀 GPU-Powered Text Generator")
+    with gr.Row():
+        inp = gr.Textbox(label="Your Prompt", placeholder="Type here...")
+    with gr.Row():
+        out = gr.Textbox(label="Generated Text", lines=5)
+    with gr.Row():
+        btn = gr.Button("Generate", variant="primary")
+    
+    btn.click(fn=generate_text, inputs=inp, outputs=out)
+    inp.submit(fn=generate_text, inputs=inp, outputs=out)
 
-demo.launch()
+# 7. Launch with GPU monitoring
+if __name__ == "__main__":
+    print(f"GPU Memory Allocated: {torch.cuda.memory_allocated()/1e9:.2f} GB")
+    demo.launch(server_name="0.0.0.0")