Update app.py

app.py

@@ -1,48 +1,30 @@
 import gradio as gr
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-MODEL_NAME = "darkc0de/XortronCriminalComputingConfig"
+# load the Xortron model
+MODEL_ID = "darkc0de/XortronCriminalComputingConfig"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
 
-print(f"Loading model: {MODEL_NAME}")
-
-# Load tokenizer & model
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-
-# 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,
+    inputs = tokenizer(message, return_tensors="pt")
+    outputs = model.generate(
+        **inputs,
         max_new_tokens=256,
         do_sample=True,
         temperature=0.7,
-        truncation=True,
-    )[0]["generated_text"]
-
-    for i in range(0, len(output), 20):
-        yield {"role": "assistant", "content": output[: i + 20]}
+        top_p=0.9,
+    )
+    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return reply
 
-# Build Gradio chat
-chat = gr.ChatInterface(
+demo = gr.ChatInterface(
     fn=respond,
-    type="messages",
-    chatbot=gr.Chatbot(height=600, show_copy_button=True, type="messages"),
+    type="messages",  # avoids that deprecation warning
+    chatbot=gr.Chatbot(height=600, show_copy_button=True),
+    textbox=gr.Textbox(placeholder="Chat with Xortron...", container=False, scale=7),
+    title="Xortron Chat",
 )
 
 if __name__ == "__main__":
-    chat.launch()
+    demo.launch()