Update app.py

app.py

@@ -1,36 +1,43 @@
 import gradio as gr
-from transformers import pipeline
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-# Auto-select model based on device
-if torch.cuda.is_available():
 MODEL_NAME = "darkc0de/XortronCriminalComputingConfig"
-    device = 0
-else:
-    MODEL_NAME = "gpt2"  # CPU fallback
-    device = -1
 
-print(f"Loading model: {MODEL_NAME} on {'GPU' if device == 0 else 'CPU'}")
+print(f"Loading model: {MODEL_NAME}")
 
-# Load Hugging Face pipeline
-generator = pipeline("text-generation", model=MODEL_NAME, device=device)
+# Load tokenizer & model
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-# Streaming response function
+# device_map="auto" lets it use GPU if available, otherwise CPU (warning: very slow on CPU)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    device_map="auto",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    low_cpu_mem_usage=True,
+)
+
+generator = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=0 if torch.cuda.is_available() else -1,
+)
+
+# Streaming response
 def respond(message, history):
     output = generator(
         message,
-        max_new_tokens=256,   # use this instead of max_length
-        num_return_sequences=1,
+        max_new_tokens=256,
         do_sample=True,
         temperature=0.7,
-        truncation=True,      # fixes truncation warning
+        truncation=True,
     )[0]["generated_text"]
 
-    # Stream output in chunks
     for i in range(0, len(output), 20):
         yield {"role": "assistant", "content": output[: i + 20]}
 
-# Build the Gradio chat
+# Build Gradio chat
 chat = gr.ChatInterface(
     fn=respond,
     type="messages",