Update app.py

app.py

@@ -1,4 +1,4 @@
-import os, re, time, datetime, traceback
+import os, re, time, datetime, traceback, torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from transformers.utils import logging as hf_logging
@@ -25,15 +25,13 @@ def log(msg: str):
 # Configuration
 # ---------------------------------------------------------------------------
 MODEL_ID = "ibm-granite/granite-3.3-2b-instruct"
-MAX_TURNS = 4               # retain last N user/AI exchanges
-MAX_TOKENS = 64
-MAX_INPUT_CH = 300
+MAX_TURNS, MAX_TOKENS, MAX_INPUT_CH = 4, 64, 300
 
 SYSTEM_MSG = (
-    "You are **SchoolSpirit AI**, the digital mascot for SchoolSpirit AI LLC, "
+    "You are **SchoolSpirit AI**, the digital mascot for SchoolSpirit AI LLC, "
     "founded by Charles Norton in 2025. The company installs on‑prem AI chat "
     "mascots, offers custom fine‑tuning of language models, and ships turnkey "
-    "GPU hardware to K‑12 schools.\n\n"
+    "PC's with preinstalled language models to K‑12 schools.\n\n"
     "GUIDELINES:\n"
     "• Use a warm, encouraging tone fit for students, parents, and staff.\n"
     "• Keep replies short—no more than four sentences unless asked.\n"
@@ -41,19 +39,32 @@ SYSTEM_MSG = (
     "• Never collect personal data or provide medical, legal, or financial advice.\n"
     "• No profanity, politics, or mature themes."
 )
-
 WELCOME_MSG = "Welcome to SchoolSpirit AI! Do you have any questions?"
 
 # ---------------------------------------------------------------------------
-# Load model
+# Load model (GPU FP‑16 if available → CPU fallback)
 # ---------------------------------------------------------------------------
 hf_logging.set_verbosity_error()
 try:
-    log("Loading model …")
+    log("Loading tokenizer …")
     tok = AutoTokenizer.from_pretrained(MODEL_ID)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID, device_map="auto", torch_dtype="auto"
-    )
+
+    if torch.cuda.is_available():
+        log("GPU detected → loading model in FP‑16")
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            device_map="auto",          # put layers on available GPU(s)
+            torch_dtype=torch.float16,
+        )
+    else:
+        log("No GPU → loading model on CPU (FP‑32)")
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            device_map="cpu",
+            torch_dtype="auto",
+            low_cpu_mem_usage=True,
+        )
+
     gen = pipeline(
         "text-generation",
         model=model,
@@ -77,15 +88,8 @@ trim = lambda m: m if len(m) <= 1 + MAX_TURNS * 2 else [m[0]] + m[-MAX_TURNS * 2
 
 
 def chat_fn(user_msg: str, history: list):
-    """
-    Gradio passes:
-      user_msg : str
-      history  : list[dict]  -> [{'role':'assistant'|'user','content':...}, ...]
-    Return a string; ChatInterface will append it as assistant message.
-    """
     log(f"User sent {len(user_msg)} chars")
 
-    # Inject system message once
     if not history or history[0]["role"] != "system":
         history.insert(0, {"role": "system", "content": SYSTEM_MSG})
 
@@ -131,7 +135,7 @@ gr.ChatInterface(
     chatbot=gr.Chatbot(
         height=480,
         type="messages",
-        value=[{"role": "assistant", "content": WELCOME_MSG}],  # preloaded welcome
+        value=[{"role": "assistant", "content": WELCOME_MSG}],
     ),
     title="SchoolSpirit AI Chat",
     theme=gr.themes.Soft(primary_hue="blue"),