Update app.py

app.py

@@ -9,24 +9,41 @@ from html import escape
 PRIMARY_MODEL = "Smilyai-labs/Sam-reason-A3"
 FALLBACK_MODEL = "Smilyai-labs/Sam-reason-A1"
 USAGE_LIMIT = 5
-RESET_AFTER_SECONDS = 20 * 60  # 20 minutes
+RESET_AFTER_SECONDS = 20 * 60
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Globals
 primary_model = primary_tokenizer = None
 fallback_model = fallback_tokenizer = None
-usage_info = {}  # { ip: {"count": int, "last_seen": float} }
+usage_info = {}
 
-# Load both models
+# Load models
 def load_models():
     global primary_model, primary_tokenizer, fallback_model, fallback_tokenizer
-    primary_tokenizer = AutoTokenizer.from_pretrained(PRIMARY_MODEL)
-    primary_model = AutoModelForCausalLM.from_pretrained(PRIMARY_MODEL).to(device).eval()
-    fallback_tokenizer = AutoTokenizer.from_pretrained(FALLBACK_MODEL)
-    fallback_model = AutoModelForCausalLM.from_pretrained(FALLBACK_MODEL).to(device).eval()
+    primary_tokenizer = AutoTokenizer.from_pretrained(PRIMARY_MODEL, trust_remote_code=True)
+    primary_model = AutoModelForCausalLM.from_pretrained(PRIMARY_MODEL, torch_dtype=torch.float16).to(device).eval()
+    fallback_tokenizer = AutoTokenizer.from_pretrained(FALLBACK_MODEL, trust_remote_code=True)
+    fallback_model = AutoModelForCausalLM.from_pretrained(FALLBACK_MODEL, torch_dtype=torch.float16).to(device).eval()
     return f"✅ Loaded: {PRIMARY_MODEL} with fallback {FALLBACK_MODEL}"
 
-# Token streaming
+# Format multi-turn history
+def build_chat_prompt(history, user_input, reasoning_enabled):
+    prefix = "/think " if reasoning_enabled else "/no_think "
+    prompt = ""
+    for user_msg, bot_msg in history:
+        prompt += f"<|user|>\n{user_msg}\n<|assistant|>\n{bot_msg}\n"
+    prompt += f"<|user|>\n{user_input}\n<|assistant|>\n"
+    return prefix + prompt
+
+# Collapse <think> block
+def format_thinking(text):
+    match = re.search(r"<think>(.*?)</think>", text, re.DOTALL)
+    if match:
+        reasoning = escape(match.group(1).strip())
+        visible = re.sub(r"<think>.*?</think>", "[thinking...]", text, flags=re.DOTALL).strip()
+        return f"{escape(visible)}<br><details><summary>🧠 Show reasoning</summary><pre>{reasoning}</pre></details>"
+    return escape(text)
+
+# Token stream generator
 def generate_stream(prompt, use_fallback=False, max_length=100, temperature=0.2, top_p=0.9):
     model = fallback_model if use_fallback else primary_model
     tokenizer = fallback_tokenizer if use_fallback else primary_tokenizer
@@ -39,7 +56,7 @@ def generate_stream(prompt, use_fallback=False, max_length=100, temperature=0.2,
         sorted_logits, sorted_indices = torch.sort(logits, descending=True)
         probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
         mask = probs > top_p
-        mask[..., 1:] = mask[..., :-1].clone()  # clone to avoid in-place issues
+        mask[..., 1:] = mask[..., :-1].clone()
         mask[..., 0] = 0
         filtered = logits.clone()
         filtered[:, sorted_indices[mask]] = -float("Inf")
@@ -51,22 +68,12 @@ def generate_stream(prompt, use_fallback=False, max_length=100, temperature=0.2,
         if next_token.item() == tokenizer.eos_token_id:
             break
 
-# Collapse <think> blocks
-def format_thinking(text):
-    match = re.search(r"<think>(.*?)</think>", text, re.DOTALL)
-    if match:
-        reasoning = escape(match.group(1).strip())
-        visible = re.sub(r"<think>.*?</think>", "[thinking...]", text, flags=re.DOTALL).strip()
-        return f"{escape(visible)}<br><details><summary>🧠 Show reasoning</summary><pre>{reasoning}</pre></details>"
-    return escape(text)
-
-# Handle each message
+# Response pipeline
 def respond(message, history, reasoning_enabled, request: gr.Request):
     ip = request.client.host if request else "unknown"
     now = time.time()
     info = usage_info.get(ip, {"count": 0, "last_seen": 0})
 
-    # Reset after idle period
     if now - info["last_seen"] > RESET_AFTER_SECONDS:
         info["count"] = 0
 
@@ -76,10 +83,9 @@ def respond(message, history, reasoning_enabled, request: gr.Request):
 
     use_fallback = info["count"] > USAGE_LIMIT
     remaining = max(0, USAGE_LIMIT - info["count"])
-    model_used = "A1" if not use_fallback else "Fallback S2.1"
+    model_used = "A3" if not use_fallback else "Fallback A1"
 
-    prefix = "/think " if reasoning_enabled else "/no_think "
-    prompt = prefix + message.strip()
+    prompt = build_chat_prompt(history, message.strip(), reasoning_enabled)
     history = history + [[message, ""]]
 
     for output in generate_stream(prompt, use_fallback=use_fallback):
@@ -92,7 +98,7 @@ def clear_chat():
 
 # UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 SamAI – Reasoning Chat")
+    gr.Markdown("# 🤖 SamAI – Reasoning Chat (Chat Mode Enabled)")
     model_status = gr.Textbox(interactive=False, label="Model Status")
     usage_counter = gr.Textbox(value="🧠 A3 messages left: 5", interactive=False, show_label=False)
     chat_box = gr.Chatbot(type="tuples")