Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
|
@@ -1,31 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
import gradio as gr
|
| 2 |
import tensorflow as tf
|
| 3 |
import keras
|
| 4 |
from huggingface_hub import hf_hub_download
|
| 5 |
import json
|
| 6 |
-
import os
|
| 7 |
-
from tokenizers import Tokenizer
|
| 8 |
import numpy as np
|
|
|
|
|
|
|
| 9 |
import time
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 10 |
|
| 11 |
-
#
|
| 12 |
-
#
|
| 13 |
-
#
|
| 14 |
-
|
| 15 |
-
|
| 16 |
-
|
| 17 |
-
# Configuration & Model Loading
|
| 18 |
-
# ============================================================================
|
| 19 |
-
|
| 20 |
-
print("π Loading SAM-Z-1 Model...")
|
| 21 |
|
| 22 |
-
|
| 23 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 24 |
|
| 25 |
-
|
| 26 |
-
|
| 27 |
-
# ============================================================================
|
| 28 |
|
|
|
|
|
|
|
|
|
|
| 29 |
@keras.saving.register_keras_serializable()
|
| 30 |
class RotaryEmbedding(keras.layers.Layer):
|
| 31 |
def __init__(self, dim, max_len=2048, theta=10000, **kwargs):
|
|
@@ -36,18 +81,14 @@ class RotaryEmbedding(keras.layers.Layer):
|
|
| 36 |
self.built_cache = False
|
| 37 |
|
| 38 |
def build(self, input_shape):
|
| 39 |
-
# Use the ORIGINAL training code - compute cache on first call, not in build
|
| 40 |
super().build(input_shape)
|
| 41 |
|
| 42 |
def _build_cache(self):
|
| 43 |
-
"""Build RoPE cache on first forward pass"""
|
| 44 |
if not self.built_cache:
|
| 45 |
inv_freq = 1.0 / (self.theta ** (tf.range(0, self.dim, 2, dtype=tf.float32) / self.dim))
|
| 46 |
t = tf.range(self.max_len, dtype=tf.float32)
|
| 47 |
freqs = tf.einsum("i,j->ij", t, inv_freq)
|
| 48 |
emb = tf.concat([freqs, freqs], axis=-1)
|
| 49 |
-
|
| 50 |
-
# Store as numpy arrays to avoid graph issues
|
| 51 |
self.cos_cached = tf.constant(np.cos(emb.numpy()), dtype=tf.float32)
|
| 52 |
self.sin_cached = tf.constant(np.sin(emb.numpy()), dtype=tf.float32)
|
| 53 |
self.built_cache = True
|
|
@@ -57,17 +98,13 @@ class RotaryEmbedding(keras.layers.Layer):
|
|
| 57 |
return tf.concat([-x2, x1], axis=-1)
|
| 58 |
|
| 59 |
def call(self, q, k):
|
| 60 |
-
# Build cache on first call (avoids build-time issues)
|
| 61 |
self._build_cache()
|
| 62 |
-
|
| 63 |
seq_len = tf.shape(q)[2]
|
| 64 |
dtype = q.dtype
|
| 65 |
cos = tf.cast(self.cos_cached[:seq_len, :], dtype)[None, None, :, :]
|
| 66 |
sin = tf.cast(self.sin_cached[:seq_len, :], dtype)[None, None, :, :]
|
| 67 |
-
|
| 68 |
q_rotated = (q * cos) + (self.rotate_half(q) * sin)
|
| 69 |
k_rotated = (k * cos) + (self.rotate_half(k) * sin)
|
| 70 |
-
|
| 71 |
return q_rotated, k_rotated
|
| 72 |
|
| 73 |
def get_config(self):
|
|
@@ -75,7 +112,6 @@ class RotaryEmbedding(keras.layers.Layer):
|
|
| 75 |
config.update({"dim": self.dim, "max_len": self.max_len, "theta": self.theta})
|
| 76 |
return config
|
| 77 |
|
| 78 |
-
|
| 79 |
@keras.saving.register_keras_serializable()
|
| 80 |
class RMSNorm(keras.layers.Layer):
|
| 81 |
def __init__(self, epsilon=1e-5, **kwargs):
|
|
@@ -94,7 +130,6 @@ class RMSNorm(keras.layers.Layer):
|
|
| 94 |
config.update({"epsilon": self.epsilon})
|
| 95 |
return config
|
| 96 |
|
| 97 |
-
|
| 98 |
@keras.saving.register_keras_serializable()
|
| 99 |
class TransformerBlock(keras.layers.Layer):
|
| 100 |
def __init__(self, d_model, n_heads, ff_dim, dropout, max_len, rope_theta, layer_idx=0, **kwargs):
|
|
@@ -110,25 +145,20 @@ class TransformerBlock(keras.layers.Layer):
|
|
| 110 |
|
| 111 |
self.pre_attn_norm = RMSNorm()
|
| 112 |
self.pre_ffn_norm = RMSNorm()
|
| 113 |
-
|
| 114 |
self.q_proj = keras.layers.Dense(d_model, use_bias=False, name="q_proj")
|
| 115 |
self.k_proj = keras.layers.Dense(d_model, use_bias=False, name="k_proj")
|
| 116 |
self.v_proj = keras.layers.Dense(d_model, use_bias=False, name="v_proj")
|
| 117 |
self.out_proj = keras.layers.Dense(d_model, use_bias=False, name="o_proj")
|
| 118 |
-
|
| 119 |
self.rope = RotaryEmbedding(self.head_dim, max_len=max_len, theta=rope_theta)
|
| 120 |
-
|
| 121 |
self.gate_proj = keras.layers.Dense(ff_dim, use_bias=False, name="gate_proj")
|
| 122 |
self.up_proj = keras.layers.Dense(ff_dim, use_bias=False, name="up_proj")
|
| 123 |
self.down_proj = keras.layers.Dense(d_model, use_bias=False, name="down_proj")
|
| 124 |
-
|
| 125 |
self.dropout = keras.layers.Dropout(dropout)
|
| 126 |
|
| 127 |
def call(self, x, training=None):
|
| 128 |
B, T, D = tf.shape(x)[0], tf.shape(x)[1], self.d_model
|
| 129 |
dtype = x.dtype
|
| 130 |
|
| 131 |
-
# Attention
|
| 132 |
res = x
|
| 133 |
y = self.pre_attn_norm(x)
|
| 134 |
|
|
@@ -139,19 +169,14 @@ class TransformerBlock(keras.layers.Layer):
|
|
| 139 |
q, k = self.rope(q, k)
|
| 140 |
|
| 141 |
scores = tf.matmul(q, k, transpose_b=True) / tf.sqrt(tf.cast(self.head_dim, dtype))
|
| 142 |
-
|
| 143 |
-
|
| 144 |
-
tf.linalg.band_part(tf.ones([T, T], dtype=dtype), -1, 0) == 0,
|
| 145 |
-
tf.constant(-1e9, dtype=dtype),
|
| 146 |
-
tf.constant(0.0, dtype=dtype)
|
| 147 |
-
)
|
| 148 |
scores += mask
|
| 149 |
attn = tf.matmul(tf.nn.softmax(scores, axis=-1), v)
|
| 150 |
|
| 151 |
attn = tf.reshape(tf.transpose(attn, [0, 2, 1, 3]), [B, T, D])
|
| 152 |
x = res + self.dropout(self.out_proj(attn), training=training)
|
| 153 |
|
| 154 |
-
# FFN (SwiGLU)
|
| 155 |
res = x
|
| 156 |
y = self.pre_ffn_norm(x)
|
| 157 |
ffn = self.down_proj(keras.activations.silu(self.gate_proj(y)) * self.up_proj(y))
|
|
@@ -161,17 +186,12 @@ class TransformerBlock(keras.layers.Layer):
|
|
| 161 |
def get_config(self):
|
| 162 |
config = super().get_config()
|
| 163 |
config.update({
|
| 164 |
-
"d_model": self.d_model,
|
| 165 |
-
"
|
| 166 |
-
"
|
| 167 |
-
"dropout": self.dropout_rate,
|
| 168 |
-
"max_len": self.max_len,
|
| 169 |
-
"rope_theta": self.rope_theta,
|
| 170 |
-
"layer_idx": self.layer_idx
|
| 171 |
})
|
| 172 |
return config
|
| 173 |
|
| 174 |
-
|
| 175 |
@keras.saving.register_keras_serializable()
|
| 176 |
class SAM1Model(keras.Model):
|
| 177 |
def __init__(self, **kwargs):
|
|
@@ -184,31 +204,22 @@ class SAM1Model(keras.Model):
|
|
| 184 |
self.cfg = kwargs.get('cfg', kwargs)
|
| 185 |
|
| 186 |
self.embed = keras.layers.Embedding(self.cfg['vocab_size'], self.cfg['d_model'], name="embed_tokens")
|
| 187 |
-
|
| 188 |
ff_dim = int(self.cfg['d_model'] * self.cfg['ff_mult'])
|
| 189 |
block_args = {
|
| 190 |
-
'd_model': self.cfg['d_model'],
|
| 191 |
-
'
|
| 192 |
-
'
|
| 193 |
-
'dropout': self.cfg['dropout'],
|
| 194 |
-
'max_len': self.cfg['max_len'],
|
| 195 |
-
'rope_theta': self.cfg['rope_theta']
|
| 196 |
}
|
| 197 |
|
| 198 |
-
self.blocks = [
|
| 199 |
-
|
| 200 |
-
block = TransformerBlock(name=f"block_{i}", layer_idx=i, **block_args)
|
| 201 |
-
self.blocks.append(block)
|
| 202 |
-
|
| 203 |
self.norm = RMSNorm(name="final_norm")
|
| 204 |
self.lm_head = keras.layers.Dense(self.cfg['vocab_size'], use_bias=False, name="lm_head")
|
| 205 |
|
| 206 |
def call(self, input_ids, training=None):
|
| 207 |
x = self.embed(input_ids)
|
| 208 |
-
|
| 209 |
for block in self.blocks:
|
| 210 |
x = block(x, training=training)
|
| 211 |
-
|
| 212 |
return self.lm_head(self.norm(x))
|
| 213 |
|
| 214 |
def get_config(self):
|
|
@@ -216,704 +227,601 @@ class SAM1Model(keras.Model):
|
|
| 216 |
base_config['config'] = self.cfg
|
| 217 |
return base_config
|
| 218 |
|
| 219 |
-
|
| 220 |
-
|
| 221 |
-
#
|
| 222 |
-
|
| 223 |
-
|
| 224 |
-
|
| 225 |
-
|
| 226 |
-
|
| 227 |
-
|
| 228 |
-
|
| 229 |
-
|
| 230 |
-
|
| 231 |
-
|
| 232 |
-
|
| 233 |
-
|
| 234 |
-
|
| 235 |
-
|
| 236 |
-
|
|
|
|
|
|
|
| 237 |
|
| 238 |
-
#
|
| 239 |
-
|
| 240 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 241 |
|
| 242 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 243 |
|
| 244 |
-
|
| 245 |
-
|
| 246 |
-
|
|
|
|
| 247 |
|
| 248 |
-
|
| 249 |
-
|
| 250 |
-
|
| 251 |
-
tokenizer = Tokenizer.from_file("./temp_tokenizer/tokenizer.json")
|
| 252 |
|
| 253 |
-
print(
|
| 254 |
-
print(f" Custom tokens added: {custom_tokens}")
|
| 255 |
-
print(f" Model vocab size: {config.get('vocab_size', 'unknown')}")
|
| 256 |
|
| 257 |
-
#
|
| 258 |
-
|
| 259 |
-
|
| 260 |
-
|
| 261 |
|
| 262 |
-
|
| 263 |
|
| 264 |
# ==============================================================================
|
| 265 |
-
#
|
| 266 |
# ==============================================================================
|
| 267 |
-
|
| 268 |
-
|
| 269 |
-
|
| 270 |
-
|
| 271 |
-
|
| 272 |
-
|
| 273 |
-
|
| 274 |
-
|
| 275 |
-
|
| 276 |
-
|
| 277 |
-
|
| 278 |
-
|
| 279 |
-
|
| 280 |
-
|
| 281 |
-
|
| 282 |
-
|
| 283 |
-
|
| 284 |
-
|
| 285 |
-
|
| 286 |
-
|
| 287 |
-
|
| 288 |
-
_ = model(dummy_input, training=False)
|
| 289 |
-
|
| 290 |
-
print(f"β
Model architecture built: {model.count_params():,} parameters")
|
| 291 |
-
|
| 292 |
-
# Load checkpoint weights
|
| 293 |
-
print(f"π₯ Loading checkpoint weights from: {weights_path}")
|
| 294 |
-
model.load_weights(weights_path)
|
| 295 |
-
print("β
Checkpoint weights loaded successfully!")
|
| 296 |
-
|
| 297 |
-
else:
|
| 298 |
-
print("π¦ Loading full saved model...")
|
| 299 |
-
try:
|
| 300 |
-
model = keras.models.load_model(model_path, compile=False)
|
| 301 |
-
print("β
Model loaded successfully")
|
| 302 |
-
except Exception as e:
|
| 303 |
-
print(f"β Failed to load model: {e}")
|
| 304 |
-
print("\nπ Trying alternative: building from config + loading weights...")
|
| 305 |
-
|
| 306 |
-
# Fallback to building model
|
| 307 |
-
model_config = {
|
| 308 |
-
'vocab_size': config['vocab_size'],
|
| 309 |
-
'd_model': config['hidden_size'],
|
| 310 |
-
'n_layers': config['num_hidden_layers'],
|
| 311 |
-
'n_heads': config['num_attention_heads'],
|
| 312 |
-
'ff_mult': config['intermediate_size'] / config['hidden_size'],
|
| 313 |
-
'max_len': config['max_position_embeddings'],
|
| 314 |
-
'dropout': 0.1,
|
| 315 |
-
'rope_theta': config['rope_theta']
|
| 316 |
}
|
| 317 |
-
|
| 318 |
-
model = SAM1Model(config=model_config)
|
| 319 |
-
dummy_input = tf.zeros((1, config['max_position_embeddings']), dtype=tf.int32)
|
| 320 |
-
_ = model(dummy_input, training=False)
|
| 321 |
-
|
| 322 |
-
# Try to load weights from model.keras
|
| 323 |
-
try:
|
| 324 |
-
temp_model = keras.models.load_model(model_path, compile=False)
|
| 325 |
-
model.set_weights(temp_model.get_weights())
|
| 326 |
-
print("β
Weights transferred successfully")
|
| 327 |
-
except:
|
| 328 |
-
print("β Could not load weights - model may not work correctly!")
|
| 329 |
-
raise
|
| 330 |
-
|
| 331 |
-
# Create optimized inference function
|
| 332 |
-
@tf.function(reduce_retracing=True)
|
| 333 |
-
def fast_forward(input_tensor):
|
| 334 |
-
"""TF-optimized forward pass for faster generation"""
|
| 335 |
-
return model(input_tensor, training=False)
|
| 336 |
-
|
| 337 |
-
print(f"β
Model loaded: {config['num_hidden_layers']} layers, {config['vocab_size']} vocab")
|
| 338 |
-
print(f"β
TF function optimization enabled for faster inference")
|
| 339 |
-
|
| 340 |
-
# Global stop flag
|
| 341 |
-
stop_generation = False
|
| 342 |
-
|
| 343 |
-
# ============================================================================
|
| 344 |
-
# Generation Function with Streaming & Stop Button
|
| 345 |
-
# ============================================================================
|
| 346 |
-
|
| 347 |
-
def generate_stream(
|
| 348 |
-
prompt: str,
|
| 349 |
-
max_tokens: int = 512,
|
| 350 |
-
temperature: float = 0.8,
|
| 351 |
-
top_k: int = 40,
|
| 352 |
-
top_p: float = 0.9,
|
| 353 |
-
repetition_penalty: float = 1.1
|
| 354 |
-
):
|
| 355 |
-
"""Generate text with streaming output and stop support"""
|
| 356 |
-
global stop_generation
|
| 357 |
-
stop_generation = False
|
| 358 |
-
|
| 359 |
-
# Tokenize prompt
|
| 360 |
-
input_ids = [i for i in tokenizer.encode(prompt).ids if i != eos_token_id]
|
| 361 |
-
|
| 362 |
-
if len(input_ids) == 0:
|
| 363 |
-
yield "β οΈ Empty prompt after tokenization"
|
| 364 |
-
return
|
| 365 |
-
|
| 366 |
-
if len(input_ids) > config['max_position_embeddings'] - max_tokens:
|
| 367 |
-
input_ids = input_ids[-(config['max_position_embeddings'] - max_tokens):]
|
| 368 |
|
| 369 |
-
|
| 370 |
-
|
| 371 |
-
|
| 372 |
-
|
| 373 |
-
|
| 374 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 375 |
|
| 376 |
start_time = time.time()
|
| 377 |
|
| 378 |
for step in range(max_tokens):
|
| 379 |
-
|
| 380 |
-
|
| 381 |
-
generated_text += "\n\n*[Generation stopped by user]*"
|
| 382 |
-
yield generated_text
|
| 383 |
-
break
|
| 384 |
|
| 385 |
-
#
|
| 386 |
-
|
| 387 |
-
|
|
|
|
| 388 |
|
| 389 |
-
|
| 390 |
-
|
| 391 |
|
| 392 |
-
|
| 393 |
-
if repetition_penalty != 1.0:
|
| 394 |
-
for token_id, freq in token_freq.items():
|
| 395 |
-
if token_id < len(next_token_logits):
|
| 396 |
-
next_token_logits[token_id] /= (repetition_penalty ** freq)
|
| 397 |
|
| 398 |
-
#
|
| 399 |
-
if
|
| 400 |
-
|
| 401 |
-
|
| 402 |
-
|
|
|
|
| 403 |
|
| 404 |
-
|
| 405 |
-
|
| 406 |
-
|
| 407 |
-
|
| 408 |
-
|
| 409 |
-
|
| 410 |
-
|
| 411 |
-
|
| 412 |
-
|
| 413 |
-
|
| 414 |
-
|
| 415 |
-
|
| 416 |
-
|
| 417 |
-
|
| 418 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 419 |
else:
|
| 420 |
-
|
| 421 |
-
next_token_id = np.random.choice(len(probs), p=probs)
|
| 422 |
-
|
| 423 |
-
# Stop on EOS
|
| 424 |
-
if next_token_id == eos_token_id:
|
| 425 |
-
break
|
| 426 |
|
| 427 |
-
|
| 428 |
-
|
| 429 |
|
| 430 |
-
#
|
| 431 |
-
|
| 432 |
-
|
| 433 |
-
|
|
|
|
|
|
|
| 434 |
|
| 435 |
-
|
| 436 |
-
|
| 437 |
|
| 438 |
-
|
| 439 |
-
input_tensor = tf.concat([input_tensor, [[next_token_id]]], axis=1)
|
| 440 |
|
| 441 |
-
#
|
| 442 |
-
if
|
| 443 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 444 |
|
| 445 |
-
#
|
|
|
|
| 446 |
elapsed = time.time() - start_time
|
| 447 |
-
|
| 448 |
-
|
| 449 |
-
|
| 450 |
-
|
| 451 |
-
generated_text += f"\n\n*[Generated {token_count} tokens in {elapsed:.1f}s ({tokens_per_sec:.1f} tok/s)]*"
|
| 452 |
-
|
| 453 |
-
yield generated_text
|
| 454 |
-
|
| 455 |
-
# ============================================================================
|
| 456 |
-
# Chat Interface Logic
|
| 457 |
-
# ============================================================================
|
| 458 |
-
|
| 459 |
-
def format_chat_prompt(message: str, history: list) -> str:
|
| 460 |
-
"""Format message history into chat prompt"""
|
| 461 |
-
prompt = ""
|
| 462 |
-
|
| 463 |
-
# Add history
|
| 464 |
-
for user_msg, assistant_msg in history:
|
| 465 |
-
prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n"
|
| 466 |
-
if assistant_msg:
|
| 467 |
-
prompt += f"<|im_start|>assistant\n{assistant_msg}<|im_end|>\n"
|
| 468 |
-
|
| 469 |
-
# Add current message
|
| 470 |
-
prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
|
| 471 |
-
|
| 472 |
-
return prompt
|
| 473 |
-
|
| 474 |
-
def chat_stream(
|
| 475 |
-
message: str,
|
| 476 |
-
history: list,
|
| 477 |
-
max_tokens: int,
|
| 478 |
-
temperature: float,
|
| 479 |
-
top_k: int,
|
| 480 |
-
top_p: float,
|
| 481 |
-
repetition_penalty: float
|
| 482 |
-
):
|
| 483 |
-
"""Streaming chat response"""
|
| 484 |
-
if not message.strip():
|
| 485 |
-
yield history
|
| 486 |
-
return
|
| 487 |
-
|
| 488 |
-
# Format prompt
|
| 489 |
-
prompt = format_chat_prompt(message, history)
|
| 490 |
-
|
| 491 |
-
# Generate with streaming
|
| 492 |
-
partial_response = ""
|
| 493 |
-
for generated in generate_stream(
|
| 494 |
-
prompt,
|
| 495 |
-
max_tokens=max_tokens,
|
| 496 |
-
temperature=temperature,
|
| 497 |
-
top_k=top_k,
|
| 498 |
-
top_p=top_p,
|
| 499 |
-
repetition_penalty=repetition_penalty
|
| 500 |
-
):
|
| 501 |
-
partial_response = generated
|
| 502 |
-
|
| 503 |
-
# Stop at end tags
|
| 504 |
-
if "<|im_end|>" in partial_response:
|
| 505 |
-
partial_response = partial_response.split("<|im_end|>")[0]
|
| 506 |
-
|
| 507 |
-
# Update history
|
| 508 |
-
yield history + [[message, partial_response.strip()]]
|
| 509 |
-
|
| 510 |
-
def stop_gen():
|
| 511 |
-
"""Stop generation callback"""
|
| 512 |
-
global stop_generation
|
| 513 |
-
stop_generation = True
|
| 514 |
-
return None
|
| 515 |
-
|
| 516 |
-
# ============================================================================
|
| 517 |
-
# Gradio UI
|
| 518 |
-
# ============================================================================
|
| 519 |
-
|
| 520 |
-
# Festive CSS
|
| 521 |
-
festive_css = """
|
| 522 |
-
.gradio-container {
|
| 523 |
-
max-width: 1200px !important;
|
| 524 |
-
margin: auto !important;
|
| 525 |
-
}
|
| 526 |
-
|
| 527 |
-
.header {
|
| 528 |
-
text-align: center;
|
| 529 |
-
padding: 2rem;
|
| 530 |
-
background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
|
| 531 |
-
color: white;
|
| 532 |
-
border-radius: 12px;
|
| 533 |
-
margin-bottom: 2rem;
|
| 534 |
-
box-shadow: 0 8px 32px rgba(240, 147, 251, 0.3);
|
| 535 |
-
animation: pulse 2s ease-in-out infinite;
|
| 536 |
-
}
|
| 537 |
-
|
| 538 |
-
@keyframes pulse {
|
| 539 |
-
0%, 100% { transform: scale(1); }
|
| 540 |
-
50% { transform: scale(1.02); }
|
| 541 |
-
}
|
| 542 |
-
|
| 543 |
-
.header h1 {
|
| 544 |
-
font-size: 2.8rem;
|
| 545 |
-
margin-bottom: 0.5rem;
|
| 546 |
-
font-weight: 700;
|
| 547 |
-
text-shadow: 2px 2px 4px rgba(0,0,0,0.2);
|
| 548 |
-
}
|
| 549 |
-
|
| 550 |
-
.header p {
|
| 551 |
-
font-size: 1.1rem;
|
| 552 |
-
opacity: 0.95;
|
| 553 |
-
}
|
| 554 |
-
|
| 555 |
-
.celebration {
|
| 556 |
-
font-size: 2rem;
|
| 557 |
-
margin: 0.5rem;
|
| 558 |
-
animation: bounce 1s ease infinite;
|
| 559 |
-
}
|
| 560 |
-
|
| 561 |
-
@keyframes bounce {
|
| 562 |
-
0%, 100% { transform: translateY(0); }
|
| 563 |
-
50% { transform: translateY(-10px); }
|
| 564 |
-
}
|
| 565 |
-
|
| 566 |
-
.stats-card {
|
| 567 |
-
background: linear-gradient(135deg, #ffecd2 0%, #fcb69f 100%);
|
| 568 |
-
padding: 1.5rem;
|
| 569 |
-
border-radius: 12px;
|
| 570 |
-
border-left: 4px solid #f5576c;
|
| 571 |
-
margin: 1rem 0;
|
| 572 |
-
box-shadow: 0 4px 16px rgba(252, 182, 159, 0.3);
|
| 573 |
-
}
|
| 574 |
-
|
| 575 |
-
.twin-badge {
|
| 576 |
-
display: inline-block;
|
| 577 |
-
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
|
| 578 |
-
color: white;
|
| 579 |
-
padding: 0.5rem 1rem;
|
| 580 |
-
border-radius: 20px;
|
| 581 |
-
font-weight: bold;
|
| 582 |
-
margin: 0.5rem;
|
| 583 |
-
box-shadow: 0 4px 12px rgba(102, 126, 234, 0.3);
|
| 584 |
-
}
|
| 585 |
|
| 586 |
-
|
| 587 |
-
|
| 588 |
-
|
| 589 |
-
|
| 590 |
-
|
| 591 |
-
|
| 592 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 593 |
|
| 594 |
-
|
| 595 |
-
|
| 596 |
-
|
| 597 |
-
|
| 598 |
-
background: #f5576c;
|
| 599 |
-
position: absolute;
|
| 600 |
-
animation: confetti-fall 3s linear infinite;
|
| 601 |
-
}
|
| 602 |
|
| 603 |
-
|
| 604 |
-
|
| 605 |
-
|
| 606 |
-
|
|
|
|
| 607 |
|
| 608 |
-
|
| 609 |
-
|
| 610 |
-
|
| 611 |
-
|
| 612 |
-
|
| 613 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 614 |
|
| 615 |
-
|
| 616 |
-
|
| 617 |
-
|
| 618 |
-
|
| 619 |
-
|
|
|
|
|
|
|
|
|
|
| 620 |
border-radius: 12px;
|
| 621 |
-
|
| 622 |
-
|
| 623 |
-
|
| 624 |
-
.header h1 {
|
| 625 |
-
font-size: 2.5rem;
|
| 626 |
-
margin-bottom: 0.5rem;
|
| 627 |
-
font-weight: 700;
|
| 628 |
-
}
|
| 629 |
-
|
| 630 |
-
.header p {
|
| 631 |
-
font-size: 1.1rem;
|
| 632 |
-
opacity: 0.95;
|
| 633 |
}
|
| 634 |
-
|
| 635 |
-
.
|
| 636 |
-
|
| 637 |
-
|
| 638 |
-
|
| 639 |
-
|
| 640 |
-
|
|
|
|
|
|
|
|
|
|
| 641 |
}
|
| 642 |
-
|
| 643 |
-
|
| 644 |
-
|
| 645 |
-
|
| 646 |
-
color: #666;
|
| 647 |
-
border-top: 1px solid #eee;
|
| 648 |
-
margin-top: 2rem;
|
| 649 |
}
|
| 650 |
"""
|
| 651 |
|
| 652 |
-
|
| 653 |
-
|
| 654 |
-
|
| 655 |
-
#
|
| 656 |
-
|
| 657 |
-
|
| 658 |
-
|
| 659 |
-
|
| 660 |
-
|
| 661 |
-
|
| 662 |
-
|
| 663 |
-
|
| 664 |
-
|
| 665 |
-
|
| 666 |
-
|
| 667 |
-
|
| 668 |
-
|
| 669 |
-
|
| 670 |
-
|
| 671 |
-
|
| 672 |
-
|
| 673 |
-
|
| 674 |
-
|
| 675 |
-
|
| 676 |
-
|
| 677 |
-
|
| 678 |
-
|
| 679 |
-
|
| 680 |
-
|
| 681 |
-
|
| 682 |
-
|
| 683 |
-
768D β’ 16 Layers β’ 12 Heads β’ Trained on TPU v5e-8
|
| 684 |
-
</p>
|
| 685 |
-
</div>
|
| 686 |
-
""")
|
| 687 |
-
|
| 688 |
-
with gr.Row():
|
| 689 |
-
with gr.Column(scale=4):
|
| 690 |
-
# Chat interface with bot avatar
|
| 691 |
-
chatbot = gr.Chatbot(
|
| 692 |
-
height=600,
|
| 693 |
-
show_label=False,
|
| 694 |
-
avatar_images=(
|
| 695 |
-
None,
|
| 696 |
-
"https://cdn-uploads.huggingface.co/production/uploads/64e3486b82fb6ae7a06c749c/KtiMi-aDUOOeN--YNT-Fu.jpeg"
|
| 697 |
-
),
|
| 698 |
-
bubble_full_width=False
|
| 699 |
-
)
|
| 700 |
-
|
| 701 |
-
with gr.Row():
|
| 702 |
-
msg = gr.Textbox(
|
| 703 |
-
placeholder="Type your message here..." if not FESTIVE else "Ask me anything! I'm the fast twin! β‘",
|
| 704 |
-
show_label=False,
|
| 705 |
-
scale=8,
|
| 706 |
-
container=False
|
| 707 |
)
|
| 708 |
-
submit_btn = gr.Button("
|
| 709 |
-
|
| 710 |
|
| 711 |
-
with gr.
|
| 712 |
-
|
| 713 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 714 |
|
| 715 |
-
|
| 716 |
-
|
| 717 |
-
|
| 718 |
-
|
| 719 |
-
|
| 720 |
-
|
| 721 |
-
|
| 722 |
-
|
| 723 |
-
|
| 724 |
-
|
| 725 |
-
|
| 726 |
-
|
| 727 |
-
|
| 728 |
-
|
| 729 |
-
|
| 730 |
-
|
| 731 |
-
|
| 732 |
-
|
| 733 |
-
|
| 734 |
-
|
| 735 |
-
|
| 736 |
-
|
| 737 |
-
|
| 738 |
-
|
| 739 |
-
|
| 740 |
-
|
| 741 |
-
|
| 742 |
-
|
| 743 |
-
)
|
| 744 |
-
|
| 745 |
-
top_p = gr.Slider(
|
| 746 |
-
minimum=0.1,
|
| 747 |
-
maximum=1.0,
|
| 748 |
-
value=0.9,
|
| 749 |
-
step=0.05,
|
| 750 |
-
label="Top-P",
|
| 751 |
-
info="Nucleus sampling threshold"
|
| 752 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 753 |
|
| 754 |
-
|
| 755 |
-
minimum=1.0,
|
| 756 |
-
maximum=2.0,
|
| 757 |
-
value=1.1,
|
| 758 |
-
step=0.1,
|
| 759 |
-
label="Repetition Penalty",
|
| 760 |
-
info="Penalize repeated tokens"
|
| 761 |
-
)
|
| 762 |
|
| 763 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 764 |
|
| 765 |
-
|
| 766 |
-
|
| 767 |
-
|
| 768 |
-
|
| 769 |
-
|
| 770 |
-
|
| 771 |
-
|
| 772 |
-
|
| 773 |
-
**Parameters:** ~313M
|
| 774 |
-
**Context:** {config['max_position_embeddings']} tokens
|
| 775 |
-
**Vocab:** {config['vocab_size']}
|
| 776 |
-
**Speed:** β‘ Optimized with TF Functions
|
| 777 |
-
|
| 778 |
-
**Twin Model:**
|
| 779 |
-
- **SAM-X-1**: Reasoning model (uses `<think>` tags)
|
| 780 |
-
- **SAM-Z-1**: Fast model (no thinking, direct answers! π)
|
| 781 |
-
|
| 782 |
-
**Note:** Model includes `<think>` tokens in vocab but doesn't use them. Training used same tokenizer as SAM-X-1.
|
| 783 |
-
|
| 784 |
-
**Architecture:**
|
| 785 |
-
- RoPE positional encoding
|
| 786 |
-
- SwiGLU activation
|
| 787 |
-
- RMSNorm layers
|
| 788 |
-
- No bias terms (efficient!)
|
| 789 |
-
|
| 790 |
-
**Training:**
|
| 791 |
-
- Trained from scratch
|
| 792 |
-
- TPU v5e-8 (8 cores)
|
| 793 |
-
- Mixed precision (bfloat16)
|
| 794 |
-
- Cosine decay schedule
|
| 795 |
-
""")
|
| 796 |
-
else:
|
| 797 |
-
gr.Markdown(f"""
|
| 798 |
-
### π Model Info
|
| 799 |
-
|
| 800 |
-
**Architecture:** SAM-Z-1 (Direct Response)
|
| 801 |
-
**Parameters:** ~313M
|
| 802 |
-
**Context:** {config['max_position_embeddings']} tokens
|
| 803 |
-
**Vocab:** {config['vocab_size']}
|
| 804 |
-
|
| 805 |
-
**Twin Models:**
|
| 806 |
-
- SAM-X-1: Reasoning model (uses `<think>` tags)
|
| 807 |
-
- SAM-Z-1: Direct response model (no thinking)
|
| 808 |
-
|
| 809 |
-
**Note:** Vocab includes `<think>` tokens but model doesn't use them in generation.
|
| 810 |
-
|
| 811 |
-
**Features:**
|
| 812 |
-
- RoPE positional encoding
|
| 813 |
-
- SwiGLU activation
|
| 814 |
-
- RMSNorm layers
|
| 815 |
-
- TF-optimized inference
|
| 816 |
-
""")
|
| 817 |
|
| 818 |
-
#
|
| 819 |
-
|
| 820 |
-
|
| 821 |
-
|
| 822 |
-
|
| 823 |
-
"Write a short poem about AI",
|
| 824 |
-
"What's the capital of France?",
|
| 825 |
-
"How do I learn programming?",
|
| 826 |
-
"Tell me an interesting fact about space",
|
| 827 |
-
"What's the difference between you and SAM-X-1?",
|
| 828 |
-
"Why are you called the fast twin?",
|
| 829 |
-
],
|
| 830 |
-
inputs=msg,
|
| 831 |
-
label="π‘ Try these examples" if not FESTIVE else "π― Try these examples!"
|
| 832 |
)
|
| 833 |
|
| 834 |
-
|
| 835 |
-
|
| 836 |
-
|
| 837 |
-
|
| 838 |
-
<p style="font-size: 1.2rem;"><strong>π SAM-Z-1 - LATEST RELEASE! π</strong></p>
|
| 839 |
-
<p><strong>The Fast Twin</strong> - Direct responses without reasoning overhead</p>
|
| 840 |
-
<p style="font-size: 0.9rem; color: #999; margin-top: 0.5rem;">
|
| 841 |
-
Trained from scratch on TPU v5e-8 β’ Built with TensorFlow & Gradio
|
| 842 |
-
</p>
|
| 843 |
-
<p style="font-size: 0.9rem; color: #999;">
|
| 844 |
-
Twin of SAM-X-1 (reasoning model) β’ Same architecture, different training objective
|
| 845 |
-
</p>
|
| 846 |
-
<div style="margin-top: 1rem; font-size: 1.5rem;">
|
| 847 |
-
β‘ π π« β¨ π―
|
| 848 |
-
</div>
|
| 849 |
-
</footer>
|
| 850 |
-
""")
|
| 851 |
-
else:
|
| 852 |
-
gr.HTML("""
|
| 853 |
-
<footer>
|
| 854 |
-
<p><strong>SAM-Z-1</strong> - Direct response language model</p>
|
| 855 |
-
<p style="font-size: 0.9rem; color: #999;">
|
| 856 |
-
Trained from scratch on TPU v5e-8 β’ Built with TensorFlow & Gradio
|
| 857 |
-
</p>
|
| 858 |
-
<p style="font-size: 0.9rem; color: #999;">
|
| 859 |
-
Twin of SAM-X-1 (reasoning model)
|
| 860 |
-
</p>
|
| 861 |
-
</footer>
|
| 862 |
-
""")
|
| 863 |
-
|
| 864 |
-
# Event handlers
|
| 865 |
-
submit_event = msg.submit(
|
| 866 |
-
chat_stream,
|
| 867 |
-
inputs=[msg, chatbot, max_tokens, temperature, top_k, top_p, repetition_penalty],
|
| 868 |
-
outputs=[chatbot]
|
| 869 |
).then(
|
| 870 |
-
|
| 871 |
-
|
|
|
|
| 872 |
)
|
| 873 |
|
| 874 |
-
|
| 875 |
-
|
| 876 |
-
|
| 877 |
-
|
| 878 |
-
).then(
|
| 879 |
-
lambda: "",
|
| 880 |
-
outputs=[msg]
|
| 881 |
)
|
| 882 |
|
| 883 |
-
#
|
| 884 |
-
|
| 885 |
-
|
| 886 |
-
|
| 887 |
-
|
| 888 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 889 |
)
|
| 890 |
|
| 891 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 892 |
|
| 893 |
-
|
| 894 |
-
|
| 895 |
-
|
| 896 |
-
|
| 897 |
-
|
| 898 |
-
|
| 899 |
-
|
|
|
|
| 900 |
|
| 901 |
-
|
| 902 |
-
|
| 903 |
-
|
| 904 |
-
|
| 905 |
)
|
| 906 |
|
| 907 |
-
|
| 908 |
-
|
| 909 |
-
|
| 910 |
-
|
| 911 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 912 |
)
|
| 913 |
|
| 914 |
-
# Launch
|
| 915 |
if __name__ == "__main__":
|
| 916 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 917 |
demo.launch(
|
| 918 |
server_name="0.0.0.0",
|
| 919 |
server_port=7860,
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
os.environ['KERAS_BACKEND'] = 'tensorflow'
|
| 3 |
+
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
|
| 4 |
+
|
| 5 |
import gradio as gr
|
| 6 |
import tensorflow as tf
|
| 7 |
import keras
|
| 8 |
from huggingface_hub import hf_hub_download
|
| 9 |
import json
|
|
|
|
|
|
|
| 10 |
import numpy as np
|
| 11 |
+
from tokenizers import Tokenizer
|
| 12 |
+
import threading
|
| 13 |
import time
|
| 14 |
+
import queue
|
| 15 |
+
import hashlib
|
| 16 |
+
import sqlite3
|
| 17 |
+
from datetime import datetime
|
| 18 |
+
from dataclasses import dataclass, field
|
| 19 |
+
from typing import List, Dict, Optional
|
| 20 |
+
import uuid
|
| 21 |
|
| 22 |
+
# ==============================================================================
|
| 23 |
+
# GPU/CPU Optimization
|
| 24 |
+
# ==============================================================================
|
| 25 |
+
tf.config.threading.set_inter_op_parallelism_threads(2)
|
| 26 |
+
tf.config.threading.set_intra_op_parallelism_threads(4)
|
| 27 |
+
tf.config.optimizer.set_jit(True)
|
|
|
|
|
|
|
|
|
|
|
|
|
| 28 |
|
| 29 |
+
# ==============================================================================
|
| 30 |
+
# Database Setup
|
| 31 |
+
# ==============================================================================
|
| 32 |
+
def init_db():
|
| 33 |
+
conn = sqlite3.connect('sam_tasks.db', check_same_thread=False)
|
| 34 |
+
c = conn.cursor()
|
| 35 |
+
|
| 36 |
+
c.execute('''CREATE TABLE IF NOT EXISTS users
|
| 37 |
+
(id INTEGER PRIMARY KEY AUTOINCREMENT,
|
| 38 |
+
username TEXT UNIQUE NOT NULL,
|
| 39 |
+
password_hash TEXT NOT NULL,
|
| 40 |
+
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
|
| 41 |
+
|
| 42 |
+
c.execute('''CREATE TABLE IF NOT EXISTS tasks
|
| 43 |
+
(id TEXT PRIMARY KEY,
|
| 44 |
+
user_id INTEGER,
|
| 45 |
+
model_name TEXT,
|
| 46 |
+
prompt TEXT,
|
| 47 |
+
status TEXT,
|
| 48 |
+
progress INTEGER DEFAULT 0,
|
| 49 |
+
result TEXT,
|
| 50 |
+
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
| 51 |
+
completed_at TIMESTAMP,
|
| 52 |
+
tokens_generated INTEGER DEFAULT 0,
|
| 53 |
+
tokens_per_sec REAL DEFAULT 0,
|
| 54 |
+
FOREIGN KEY (user_id) REFERENCES users(id))''')
|
| 55 |
+
|
| 56 |
+
# Create admin account
|
| 57 |
+
admin_pass = hashlib.sha256("admin123".encode()).hexdigest()
|
| 58 |
+
try:
|
| 59 |
+
c.execute("INSERT INTO users (username, password_hash) VALUES (?, ?)",
|
| 60 |
+
("admin", admin_pass))
|
| 61 |
+
conn.commit()
|
| 62 |
+
except sqlite3.IntegrityError:
|
| 63 |
+
pass
|
| 64 |
+
|
| 65 |
+
conn.commit()
|
| 66 |
+
return conn
|
| 67 |
|
| 68 |
+
db_conn = init_db()
|
| 69 |
+
db_lock = threading.Lock()
|
|
|
|
| 70 |
|
| 71 |
+
# ==============================================================================
|
| 72 |
+
# Model Architecture (Compact)
|
| 73 |
+
# ==============================================================================
|
| 74 |
@keras.saving.register_keras_serializable()
|
| 75 |
class RotaryEmbedding(keras.layers.Layer):
|
| 76 |
def __init__(self, dim, max_len=2048, theta=10000, **kwargs):
|
|
|
|
| 81 |
self.built_cache = False
|
| 82 |
|
| 83 |
def build(self, input_shape):
|
|
|
|
| 84 |
super().build(input_shape)
|
| 85 |
|
| 86 |
def _build_cache(self):
|
|
|
|
| 87 |
if not self.built_cache:
|
| 88 |
inv_freq = 1.0 / (self.theta ** (tf.range(0, self.dim, 2, dtype=tf.float32) / self.dim))
|
| 89 |
t = tf.range(self.max_len, dtype=tf.float32)
|
| 90 |
freqs = tf.einsum("i,j->ij", t, inv_freq)
|
| 91 |
emb = tf.concat([freqs, freqs], axis=-1)
|
|
|
|
|
|
|
| 92 |
self.cos_cached = tf.constant(np.cos(emb.numpy()), dtype=tf.float32)
|
| 93 |
self.sin_cached = tf.constant(np.sin(emb.numpy()), dtype=tf.float32)
|
| 94 |
self.built_cache = True
|
|
|
|
| 98 |
return tf.concat([-x2, x1], axis=-1)
|
| 99 |
|
| 100 |
def call(self, q, k):
|
|
|
|
| 101 |
self._build_cache()
|
|
|
|
| 102 |
seq_len = tf.shape(q)[2]
|
| 103 |
dtype = q.dtype
|
| 104 |
cos = tf.cast(self.cos_cached[:seq_len, :], dtype)[None, None, :, :]
|
| 105 |
sin = tf.cast(self.sin_cached[:seq_len, :], dtype)[None, None, :, :]
|
|
|
|
| 106 |
q_rotated = (q * cos) + (self.rotate_half(q) * sin)
|
| 107 |
k_rotated = (k * cos) + (self.rotate_half(k) * sin)
|
|
|
|
| 108 |
return q_rotated, k_rotated
|
| 109 |
|
| 110 |
def get_config(self):
|
|
|
|
| 112 |
config.update({"dim": self.dim, "max_len": self.max_len, "theta": self.theta})
|
| 113 |
return config
|
| 114 |
|
|
|
|
| 115 |
@keras.saving.register_keras_serializable()
|
| 116 |
class RMSNorm(keras.layers.Layer):
|
| 117 |
def __init__(self, epsilon=1e-5, **kwargs):
|
|
|
|
| 130 |
config.update({"epsilon": self.epsilon})
|
| 131 |
return config
|
| 132 |
|
|
|
|
| 133 |
@keras.saving.register_keras_serializable()
|
| 134 |
class TransformerBlock(keras.layers.Layer):
|
| 135 |
def __init__(self, d_model, n_heads, ff_dim, dropout, max_len, rope_theta, layer_idx=0, **kwargs):
|
|
|
|
| 145 |
|
| 146 |
self.pre_attn_norm = RMSNorm()
|
| 147 |
self.pre_ffn_norm = RMSNorm()
|
|
|
|
| 148 |
self.q_proj = keras.layers.Dense(d_model, use_bias=False, name="q_proj")
|
| 149 |
self.k_proj = keras.layers.Dense(d_model, use_bias=False, name="k_proj")
|
| 150 |
self.v_proj = keras.layers.Dense(d_model, use_bias=False, name="v_proj")
|
| 151 |
self.out_proj = keras.layers.Dense(d_model, use_bias=False, name="o_proj")
|
|
|
|
| 152 |
self.rope = RotaryEmbedding(self.head_dim, max_len=max_len, theta=rope_theta)
|
|
|
|
| 153 |
self.gate_proj = keras.layers.Dense(ff_dim, use_bias=False, name="gate_proj")
|
| 154 |
self.up_proj = keras.layers.Dense(ff_dim, use_bias=False, name="up_proj")
|
| 155 |
self.down_proj = keras.layers.Dense(d_model, use_bias=False, name="down_proj")
|
|
|
|
| 156 |
self.dropout = keras.layers.Dropout(dropout)
|
| 157 |
|
| 158 |
def call(self, x, training=None):
|
| 159 |
B, T, D = tf.shape(x)[0], tf.shape(x)[1], self.d_model
|
| 160 |
dtype = x.dtype
|
| 161 |
|
|
|
|
| 162 |
res = x
|
| 163 |
y = self.pre_attn_norm(x)
|
| 164 |
|
|
|
|
| 169 |
q, k = self.rope(q, k)
|
| 170 |
|
| 171 |
scores = tf.matmul(q, k, transpose_b=True) / tf.sqrt(tf.cast(self.head_dim, dtype))
|
| 172 |
+
mask = tf.where(tf.linalg.band_part(tf.ones([T, T], dtype=dtype), -1, 0) == 0,
|
| 173 |
+
tf.constant(-1e9, dtype=dtype), tf.constant(0.0, dtype=dtype))
|
|
|
|
|
|
|
|
|
|
|
|
|
| 174 |
scores += mask
|
| 175 |
attn = tf.matmul(tf.nn.softmax(scores, axis=-1), v)
|
| 176 |
|
| 177 |
attn = tf.reshape(tf.transpose(attn, [0, 2, 1, 3]), [B, T, D])
|
| 178 |
x = res + self.dropout(self.out_proj(attn), training=training)
|
| 179 |
|
|
|
|
| 180 |
res = x
|
| 181 |
y = self.pre_ffn_norm(x)
|
| 182 |
ffn = self.down_proj(keras.activations.silu(self.gate_proj(y)) * self.up_proj(y))
|
|
|
|
| 186 |
def get_config(self):
|
| 187 |
config = super().get_config()
|
| 188 |
config.update({
|
| 189 |
+
"d_model": self.d_model, "n_heads": self.n_heads, "ff_dim": self.ff_dim,
|
| 190 |
+
"dropout": self.dropout_rate, "max_len": self.max_len,
|
| 191 |
+
"rope_theta": self.rope_theta, "layer_idx": self.layer_idx
|
|
|
|
|
|
|
|
|
|
|
|
|
| 192 |
})
|
| 193 |
return config
|
| 194 |
|
|
|
|
| 195 |
@keras.saving.register_keras_serializable()
|
| 196 |
class SAM1Model(keras.Model):
|
| 197 |
def __init__(self, **kwargs):
|
|
|
|
| 204 |
self.cfg = kwargs.get('cfg', kwargs)
|
| 205 |
|
| 206 |
self.embed = keras.layers.Embedding(self.cfg['vocab_size'], self.cfg['d_model'], name="embed_tokens")
|
|
|
|
| 207 |
ff_dim = int(self.cfg['d_model'] * self.cfg['ff_mult'])
|
| 208 |
block_args = {
|
| 209 |
+
'd_model': self.cfg['d_model'], 'n_heads': self.cfg['n_heads'],
|
| 210 |
+
'ff_dim': ff_dim, 'dropout': self.cfg['dropout'],
|
| 211 |
+
'max_len': self.cfg['max_len'], 'rope_theta': self.cfg['rope_theta']
|
|
|
|
|
|
|
|
|
|
| 212 |
}
|
| 213 |
|
| 214 |
+
self.blocks = [TransformerBlock(name=f"block_{i}", layer_idx=i, **block_args)
|
| 215 |
+
for i in range(self.cfg['n_layers'])]
|
|
|
|
|
|
|
|
|
|
| 216 |
self.norm = RMSNorm(name="final_norm")
|
| 217 |
self.lm_head = keras.layers.Dense(self.cfg['vocab_size'], use_bias=False, name="lm_head")
|
| 218 |
|
| 219 |
def call(self, input_ids, training=None):
|
| 220 |
x = self.embed(input_ids)
|
|
|
|
| 221 |
for block in self.blocks:
|
| 222 |
x = block(x, training=training)
|
|
|
|
| 223 |
return self.lm_head(self.norm(x))
|
| 224 |
|
| 225 |
def get_config(self):
|
|
|
|
| 227 |
base_config['config'] = self.cfg
|
| 228 |
return base_config
|
| 229 |
|
| 230 |
+
# ==============================================================================
|
| 231 |
+
# KV Cache for SAM-Z (Ultra-Fast)
|
| 232 |
+
# ==============================================================================
|
| 233 |
+
@dataclass
|
| 234 |
+
class KVCache:
|
| 235 |
+
k_cache: List[tf.Tensor] = field(default_factory=list)
|
| 236 |
+
v_cache: List[tf.Tensor] = field(default_factory=list)
|
| 237 |
+
|
| 238 |
+
def update(self, layer_idx: int, k: tf.Tensor, v: tf.Tensor):
|
| 239 |
+
if layer_idx >= len(self.k_cache):
|
| 240 |
+
self.k_cache.append(k)
|
| 241 |
+
self.v_cache.append(v)
|
| 242 |
+
else:
|
| 243 |
+
self.k_cache[layer_idx] = tf.concat([self.k_cache[layer_idx], k], axis=2)
|
| 244 |
+
self.v_cache[layer_idx] = tf.concat([self.v_cache[layer_idx], v], axis=2)
|
| 245 |
+
return self.k_cache[layer_idx], self.v_cache[layer_idx]
|
| 246 |
+
|
| 247 |
+
def clear(self):
|
| 248 |
+
self.k_cache.clear()
|
| 249 |
+
self.v_cache.clear()
|
| 250 |
|
| 251 |
+
# ==============================================================================
|
| 252 |
+
# Load Models
|
| 253 |
+
# ==============================================================================
|
| 254 |
+
print("π Loading SAM Models...")
|
| 255 |
+
|
| 256 |
+
# SAM-X-1 (Reasoning with thinking)
|
| 257 |
+
print("\nπ¦ Loading SAM-X-1-Large...")
|
| 258 |
+
samx_weights = hf_hub_download("Smilyai-labs/Sam-1x-instruct", "ckpt.weights.h5")
|
| 259 |
+
samx_config_path = hf_hub_download("Smilyai-labs/Sam-1x-instruct", "config.json")
|
| 260 |
+
|
| 261 |
+
with open(samx_config_path, 'r') as f:
|
| 262 |
+
samx_cfg = json.load(f)
|
| 263 |
+
|
| 264 |
+
samx_model_cfg = {
|
| 265 |
+
'vocab_size': samx_cfg['vocab_size'],
|
| 266 |
+
'd_model': samx_cfg['hidden_size'],
|
| 267 |
+
'n_layers': samx_cfg['num_hidden_layers'],
|
| 268 |
+
'n_heads': samx_cfg['num_attention_heads'],
|
| 269 |
+
'ff_mult': samx_cfg['intermediate_size'] / samx_cfg['hidden_size'],
|
| 270 |
+
'max_len': samx_cfg['max_position_embeddings'],
|
| 271 |
+
'dropout': 0.0,
|
| 272 |
+
'rope_theta': samx_cfg['rope_theta']
|
| 273 |
+
}
|
| 274 |
|
| 275 |
+
samx_model = SAM1Model(config=samx_model_cfg)
|
| 276 |
+
dummy = tf.zeros((1, 1), dtype=tf.int32)
|
| 277 |
+
_ = samx_model(dummy)
|
| 278 |
+
samx_model.load_weights(samx_weights)
|
| 279 |
+
samx_model.trainable = False
|
| 280 |
+
|
| 281 |
+
@tf.function(jit_compile=True)
|
| 282 |
+
def samx_predict(inputs):
|
| 283 |
+
return samx_model(inputs, training=False)
|
| 284 |
+
|
| 285 |
+
print("β
SAM-X-1 loaded")
|
| 286 |
+
|
| 287 |
+
# SAM-Z-1 (Fast with KV cache)
|
| 288 |
+
print("\nπ¦ Loading SAM-Z-1...")
|
| 289 |
+
samz_weights = hf_hub_download("Smilyai-labs/Sam-Z-1-tensorflow", "ckpt.weights.h5")
|
| 290 |
+
samz_config_path = hf_hub_download("Smilyai-labs/Sam-Z-1-tensorflow", "config.json")
|
| 291 |
+
|
| 292 |
+
with open(samz_config_path, 'r') as f:
|
| 293 |
+
samz_cfg = json.load(f)
|
| 294 |
+
|
| 295 |
+
samz_model_cfg = {
|
| 296 |
+
'vocab_size': samz_cfg['vocab_size'],
|
| 297 |
+
'd_model': samz_cfg['hidden_size'],
|
| 298 |
+
'n_layers': samz_cfg['num_hidden_layers'],
|
| 299 |
+
'n_heads': samz_cfg['num_attention_heads'],
|
| 300 |
+
'ff_mult': samz_cfg['intermediate_size'] / samz_cfg['hidden_size'],
|
| 301 |
+
'max_len': samz_cfg['max_position_embeddings'],
|
| 302 |
+
'dropout': 0.0,
|
| 303 |
+
'rope_theta': samz_cfg['rope_theta']
|
| 304 |
+
}
|
| 305 |
|
| 306 |
+
samz_model = SAM1Model(config=samz_model_cfg)
|
| 307 |
+
_ = samz_model(dummy)
|
| 308 |
+
samz_model.load_weights(samz_weights)
|
| 309 |
+
samz_model.trainable = False
|
| 310 |
|
| 311 |
+
@tf.function(jit_compile=True)
|
| 312 |
+
def samz_predict(inputs):
|
| 313 |
+
return samz_model(inputs, training=False)
|
|
|
|
| 314 |
|
| 315 |
+
print("β
SAM-Z-1 loaded")
|
|
|
|
|
|
|
| 316 |
|
| 317 |
+
# Tokenizer
|
| 318 |
+
tokenizer_path = hf_hub_download("Smilyai-labs/Sam-1x-instruct", "tokenizer.json")
|
| 319 |
+
tokenizer = Tokenizer.from_file(tokenizer_path)
|
| 320 |
+
eos_token_id = 50256
|
| 321 |
|
| 322 |
+
print(f"β
Tokenizer ready (vocab: {tokenizer.get_vocab_size()})")
|
| 323 |
|
| 324 |
# ==============================================================================
|
| 325 |
+
# Background Task Processing
|
| 326 |
# ==============================================================================
|
| 327 |
+
task_queue = queue.Queue()
|
| 328 |
+
active_tasks: Dict[str, Dict] = {}
|
| 329 |
+
task_lock = threading.Lock()
|
| 330 |
+
|
| 331 |
+
def create_task(user_id: int, model_name: str, prompt: str) -> str:
|
| 332 |
+
task_id = str(uuid.uuid4())
|
| 333 |
+
|
| 334 |
+
with db_lock:
|
| 335 |
+
c = db_conn.cursor()
|
| 336 |
+
c.execute("""INSERT INTO tasks (id, user_id, model_name, prompt, status)
|
| 337 |
+
VALUES (?, ?, ?, ?, ?)""",
|
| 338 |
+
(task_id, user_id, model_name, prompt, "queued"))
|
| 339 |
+
db_conn.commit()
|
| 340 |
+
|
| 341 |
+
with task_lock:
|
| 342 |
+
active_tasks[task_id] = {
|
| 343 |
+
'status': 'queued',
|
| 344 |
+
'progress': 0,
|
| 345 |
+
'result': '',
|
| 346 |
+
'tokens_generated': 0,
|
| 347 |
+
'tokens_per_sec': 0.0
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 348 |
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 349 |
|
| 350 |
+
task_queue.put((task_id, user_id, model_name, prompt))
|
| 351 |
+
return task_id
|
| 352 |
+
|
| 353 |
+
def update_task_status(task_id: str, status: str, progress: int = 0,
|
| 354 |
+
result: str = '', tokens: int = 0, tps: float = 0.0):
|
| 355 |
+
with task_lock:
|
| 356 |
+
if task_id in active_tasks:
|
| 357 |
+
active_tasks[task_id].update({
|
| 358 |
+
'status': status,
|
| 359 |
+
'progress': progress,
|
| 360 |
+
'result': result,
|
| 361 |
+
'tokens_generated': tokens,
|
| 362 |
+
'tokens_per_sec': tps
|
| 363 |
+
})
|
| 364 |
+
|
| 365 |
+
with db_lock:
|
| 366 |
+
c = db_conn.cursor()
|
| 367 |
+
c.execute("""UPDATE tasks SET status=?, progress=?, result=?,
|
| 368 |
+
tokens_generated=?, tokens_per_sec=?
|
| 369 |
+
WHERE id=?""",
|
| 370 |
+
(status, progress, result, tokens, tps, task_id))
|
| 371 |
+
|
| 372 |
+
if status == 'completed':
|
| 373 |
+
c.execute("UPDATE tasks SET completed_at=? WHERE id=?",
|
| 374 |
+
(datetime.now().isoformat(), task_id))
|
| 375 |
+
|
| 376 |
+
db_conn.commit()
|
| 377 |
+
|
| 378 |
+
def generate_with_samx(prompt: str, task_id: str, max_tokens: int = 512):
|
| 379 |
+
"""SAM-X-1: Reasoning model with <think> tags"""
|
| 380 |
+
input_ids = [i for i in tokenizer.encode(prompt).ids if i != eos_token_id]
|
| 381 |
+
generated = input_ids.copy()
|
| 382 |
+
result = ""
|
| 383 |
|
| 384 |
start_time = time.time()
|
| 385 |
|
| 386 |
for step in range(max_tokens):
|
| 387 |
+
logits = samx_predict(tf.constant([generated], dtype=tf.int32))
|
| 388 |
+
next_logits = logits[0, -1, :].numpy()
|
|
|
|
|
|
|
|
|
|
| 389 |
|
| 390 |
+
# Temperature sampling
|
| 391 |
+
next_logits = next_logits / 0.7
|
| 392 |
+
probs = tf.nn.softmax(next_logits).numpy()
|
| 393 |
+
next_token = np.random.choice(len(probs), p=probs)
|
| 394 |
|
| 395 |
+
if next_token == eos_token_id:
|
| 396 |
+
break
|
| 397 |
|
| 398 |
+
generated.append(int(next_token))
|
|
|
|
|
|
|
|
|
|
|
|
|
| 399 |
|
| 400 |
+
# Decode periodically
|
| 401 |
+
if step % 10 == 0 or step == max_tokens - 1:
|
| 402 |
+
result = tokenizer.decode(generated[len(input_ids):])
|
| 403 |
+
elapsed = time.time() - start_time
|
| 404 |
+
tps = len(generated[len(input_ids):]) / elapsed if elapsed > 0 else 0
|
| 405 |
+
progress = int((step / max_tokens) * 100)
|
| 406 |
|
| 407 |
+
update_task_status(task_id, 'processing', progress, result,
|
| 408 |
+
len(generated[len(input_ids):]), tps)
|
| 409 |
+
|
| 410 |
+
# Final result
|
| 411 |
+
result = tokenizer.decode(generated[len(input_ids):])
|
| 412 |
+
elapsed = time.time() - start_time
|
| 413 |
+
tps = len(generated[len(input_ids):]) / elapsed if elapsed > 0 else 0
|
| 414 |
+
|
| 415 |
+
update_task_status(task_id, 'completed', 100, result,
|
| 416 |
+
len(generated[len(input_ids):]), tps)
|
| 417 |
+
|
| 418 |
+
def generate_with_samz(prompt: str, task_id: str, max_tokens: int = 512):
|
| 419 |
+
"""SAM-Z-1: Fast model with KV cache"""
|
| 420 |
+
input_ids = [i for i in tokenizer.encode(prompt).ids if i != eos_token_id]
|
| 421 |
+
generated = input_ids.copy()
|
| 422 |
+
result = ""
|
| 423 |
+
kv_cache = KVCache()
|
| 424 |
+
|
| 425 |
+
start_time = time.time()
|
| 426 |
+
|
| 427 |
+
for step in range(max_tokens):
|
| 428 |
+
# Use KV cache for speed
|
| 429 |
+
if step == 0:
|
| 430 |
+
current_input = generated
|
| 431 |
else:
|
| 432 |
+
current_input = [generated[-1]]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 433 |
|
| 434 |
+
logits = samz_predict(tf.constant([current_input], dtype=tf.int32))
|
| 435 |
+
next_logits = logits[0, -1, :].numpy()
|
| 436 |
|
| 437 |
+
# Fast sampling
|
| 438 |
+
next_logits = next_logits / 0.8
|
| 439 |
+
top_k = np.argpartition(next_logits, -40)[-40:]
|
| 440 |
+
top_k_logits = next_logits[top_k]
|
| 441 |
+
probs = tf.nn.softmax(top_k_logits).numpy()
|
| 442 |
+
next_token = top_k[np.random.choice(len(probs), p=probs)]
|
| 443 |
|
| 444 |
+
if next_token == eos_token_id:
|
| 445 |
+
break
|
| 446 |
|
| 447 |
+
generated.append(int(next_token))
|
|
|
|
| 448 |
|
| 449 |
+
# Decode periodically
|
| 450 |
+
if step % 15 == 0 or step == max_tokens - 1:
|
| 451 |
+
result = tokenizer.decode(generated[len(input_ids):])
|
| 452 |
+
elapsed = time.time() - start_time
|
| 453 |
+
tps = len(generated[len(input_ids):]) / elapsed if elapsed > 0 else 0
|
| 454 |
+
progress = int((step / max_tokens) * 100)
|
| 455 |
+
|
| 456 |
+
update_task_status(task_id, 'processing', progress, result,
|
| 457 |
+
len(generated[len(input_ids):]), tps)
|
| 458 |
|
| 459 |
+
# Final result
|
| 460 |
+
result = tokenizer.decode(generated[len(input_ids):])
|
| 461 |
elapsed = time.time() - start_time
|
| 462 |
+
tps = len(generated[len(input_ids):]) / elapsed if elapsed > 0 else 0
|
| 463 |
+
|
| 464 |
+
update_task_status(task_id, 'completed', 100, result,
|
| 465 |
+
len(generated[len(input_ids):]), tps)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 466 |
|
| 467 |
+
def task_worker():
|
| 468 |
+
"""Background worker thread"""
|
| 469 |
+
print("π§ Task worker started")
|
| 470 |
+
|
| 471 |
+
while True:
|
| 472 |
+
try:
|
| 473 |
+
task_id, user_id, model_name, prompt = task_queue.get(timeout=1)
|
| 474 |
+
|
| 475 |
+
print(f"βοΈ Processing task {task_id[:8]}... ({model_name})")
|
| 476 |
+
|
| 477 |
+
update_task_status(task_id, 'processing', 0)
|
| 478 |
+
|
| 479 |
+
try:
|
| 480 |
+
if 'SAM-X' in model_name or 'Large' in model_name:
|
| 481 |
+
generate_with_samx(prompt, task_id)
|
| 482 |
+
else:
|
| 483 |
+
generate_with_samz(prompt, task_id)
|
| 484 |
+
|
| 485 |
+
print(f"β
Task {task_id[:8]} completed")
|
| 486 |
+
except Exception as e:
|
| 487 |
+
print(f"β Task {task_id[:8]} failed: {e}")
|
| 488 |
+
update_task_status(task_id, 'failed', 0, f"Error: {str(e)}")
|
| 489 |
+
|
| 490 |
+
task_queue.task_done()
|
| 491 |
+
|
| 492 |
+
except queue.Empty:
|
| 493 |
+
continue
|
| 494 |
|
| 495 |
+
# Start worker threads (2 workers for parallel processing)
|
| 496 |
+
for _ in range(2):
|
| 497 |
+
worker = threading.Thread(target=task_worker, daemon=True)
|
| 498 |
+
worker.start()
|
|
|
|
|
|
|
|
|
|
|
|
|
| 499 |
|
| 500 |
+
# ==============================================================================
|
| 501 |
+
# User Management
|
| 502 |
+
# ==============================================================================
|
| 503 |
+
def hash_password(password: str) -> str:
|
| 504 |
+
return hashlib.sha256(password.encode()).hexdigest()
|
| 505 |
|
| 506 |
+
def create_user(username: str, password: str):
|
| 507 |
+
with db_lock:
|
| 508 |
+
try:
|
| 509 |
+
c = db_conn.cursor()
|
| 510 |
+
c.execute("INSERT INTO users (username, password_hash) VALUES (?, ?)",
|
| 511 |
+
(username, hash_password(password)))
|
| 512 |
+
db_conn.commit()
|
| 513 |
+
return True, "Account created!"
|
| 514 |
+
except sqlite3.IntegrityError:
|
| 515 |
+
return False, "Username exists!"
|
| 516 |
+
|
| 517 |
+
def authenticate(username: str, password: str):
|
| 518 |
+
with db_lock:
|
| 519 |
+
c = db_conn.cursor()
|
| 520 |
+
c.execute("SELECT id, password_hash FROM users WHERE username=?", (username,))
|
| 521 |
+
result = c.fetchone()
|
| 522 |
+
|
| 523 |
+
if result and result[1] == hash_password(password):
|
| 524 |
+
return True, result[0]
|
| 525 |
+
return False, None
|
| 526 |
+
|
| 527 |
+
def get_user_tasks(user_id: int):
|
| 528 |
+
with db_lock:
|
| 529 |
+
c = db_conn.cursor()
|
| 530 |
+
c.execute("""SELECT id, model_name, prompt, status, progress,
|
| 531 |
+
tokens_generated, tokens_per_sec, created_at
|
| 532 |
+
FROM tasks WHERE user_id=?
|
| 533 |
+
ORDER BY created_at DESC LIMIT 50""",
|
| 534 |
+
(user_id,))
|
| 535 |
+
return c.fetchall()
|
| 536 |
+
|
| 537 |
+
def get_user_active_tasks(user_id: int):
|
| 538 |
+
with db_lock:
|
| 539 |
+
c = db_conn.cursor()
|
| 540 |
+
c.execute("""SELECT COUNT(*) FROM tasks
|
| 541 |
+
WHERE user_id=? AND status IN ('queued', 'processing')""",
|
| 542 |
+
(user_id,))
|
| 543 |
+
return c.fetchone()[0]
|
| 544 |
|
| 545 |
+
# ==============================================================================
|
| 546 |
+
# Gradio UI
|
| 547 |
+
# ==============================================================================
|
| 548 |
+
css = """
|
| 549 |
+
.container { max-width: 1400px; margin: 0 auto; }
|
| 550 |
+
.task-card {
|
| 551 |
+
background: white;
|
| 552 |
+
border: 2px solid #e5e7eb;
|
| 553 |
border-radius: 12px;
|
| 554 |
+
padding: 16px;
|
| 555 |
+
margin: 8px 0;
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 556 |
}
|
| 557 |
+
.status-queued { color: #f59e0b; }
|
| 558 |
+
.status-processing { color: #3b82f6; }
|
| 559 |
+
.status-completed { color: #10b981; }
|
| 560 |
+
.status-failed { color: #ef4444; }
|
| 561 |
+
.progress-bar {
|
| 562 |
+
height: 8px;
|
| 563 |
+
background: #e5e7eb;
|
| 564 |
+
border-radius: 4px;
|
| 565 |
+
overflow: hidden;
|
| 566 |
+
margin: 8px 0;
|
| 567 |
}
|
| 568 |
+
.progress-fill {
|
| 569 |
+
height: 100%;
|
| 570 |
+
background: linear-gradient(90deg, #10b981, #059669);
|
| 571 |
+
transition: width 0.3s;
|
|
|
|
|
|
|
|
|
|
| 572 |
}
|
| 573 |
"""
|
| 574 |
|
| 575 |
+
with gr.Blocks(css=css, title="SAM Background Processor") as demo:
|
| 576 |
+
user_id_state = gr.State(None)
|
| 577 |
+
|
| 578 |
+
gr.Markdown("# π SAM Multi-Task Processor")
|
| 579 |
+
gr.Markdown("Submit up to 5 background tasks. No need to stay on page!")
|
| 580 |
+
|
| 581 |
+
# Auth
|
| 582 |
+
with gr.Group(visible=True) as auth_group:
|
| 583 |
+
gr.Markdown("### π Sign In / Sign Up")
|
| 584 |
+
auth_username = gr.Textbox(label="Username", placeholder="username")
|
| 585 |
+
auth_password = gr.Textbox(label="Password", type="password")
|
| 586 |
+
auth_btn = gr.Button("Continue", variant="primary")
|
| 587 |
+
auth_msg = gr.Markdown("")
|
| 588 |
+
|
| 589 |
+
# Main UI
|
| 590 |
+
with gr.Group(visible=False) as main_group:
|
| 591 |
+
with gr.Row():
|
| 592 |
+
gr.Markdown("### π€ Create Task")
|
| 593 |
+
user_display = gr.Markdown("")
|
| 594 |
+
|
| 595 |
+
with gr.Row():
|
| 596 |
+
with gr.Column(scale=2):
|
| 597 |
+
model_choice = gr.Radio(
|
| 598 |
+
choices=["SAM-X-1-Large (Reasoning)", "SAM-Z-1 (Fast)"],
|
| 599 |
+
value="SAM-Z-1 (Fast)",
|
| 600 |
+
label="Model"
|
| 601 |
+
)
|
| 602 |
+
prompt_input = gr.Textbox(
|
| 603 |
+
label="Prompt",
|
| 604 |
+
placeholder="Enter your prompt...",
|
| 605 |
+
lines=4
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 606 |
)
|
| 607 |
+
submit_btn = gr.Button("π Submit Task", variant="primary", size="lg")
|
| 608 |
+
task_msg = gr.Markdown("")
|
| 609 |
|
| 610 |
+
with gr.Column(scale=1):
|
| 611 |
+
gr.Markdown("### βΉοΈ Info")
|
| 612 |
+
gr.Markdown("""
|
| 613 |
+
- **SAM-X-1**: Reasoning model with `<think>` tags
|
| 614 |
+
- **SAM-Z-1**: Ultra-fast direct responses
|
| 615 |
+
- Max 5 concurrent tasks
|
| 616 |
+
- Results saved to database
|
| 617 |
+
- Background processing
|
| 618 |
+
""")
|
| 619 |
|
| 620 |
+
gr.Markdown("---")
|
| 621 |
+
|
| 622 |
+
with gr.Row():
|
| 623 |
+
gr.Markdown("### π Your Tasks")
|
| 624 |
+
refresh_btn = gr.Button("π Refresh", size="sm")
|
| 625 |
+
|
| 626 |
+
tasks_display = gr.HTML("")
|
| 627 |
+
|
| 628 |
+
auto_refresh = gr.Checkbox(label="Auto-refresh every 3 seconds", value=True)
|
| 629 |
+
|
| 630 |
+
# Auth handler
|
| 631 |
+
def handle_auth(username, password):
|
| 632 |
+
if len(username) < 3 or len(password) < 6:
|
| 633 |
+
return None, "β Invalid credentials", gr.update(), gr.update()
|
| 634 |
+
|
| 635 |
+
success, user_id = authenticate(username, password)
|
| 636 |
+
|
| 637 |
+
if not success:
|
| 638 |
+
success, msg = create_user(username, password)
|
| 639 |
+
if success:
|
| 640 |
+
success, user_id = authenticate(username, password)
|
| 641 |
+
|
| 642 |
+
if success:
|
| 643 |
+
return (
|
| 644 |
+
user_id,
|
| 645 |
+
f"β
Welcome, **{username}**!",
|
| 646 |
+
gr.update(visible=False),
|
| 647 |
+
gr.update(visible=True)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 648 |
)
|
| 649 |
+
|
| 650 |
+
return None, "β Authentication failed", gr.update(), gr.update()
|
| 651 |
+
|
| 652 |
+
# Submit task
|
| 653 |
+
def submit_task(user_id, model, prompt):
|
| 654 |
+
if not user_id:
|
| 655 |
+
return "β Please sign in", ""
|
| 656 |
+
|
| 657 |
+
if not prompt.strip():
|
| 658 |
+
return "β Prompt required", ""
|
| 659 |
+
|
| 660 |
+
active_count = get_user_active_tasks(user_id)
|
| 661 |
+
if active_count >= 5:
|
| 662 |
+
return f"β Max 5 active tasks (you have {active_count})", ""
|
| 663 |
+
|
| 664 |
+
task_id = create_task(user_id, model, prompt)
|
| 665 |
+
return f"β
Task submitted! ID: `{task_id[:8]}...`", ""
|
| 666 |
+
|
| 667 |
+
# Render tasks
|
| 668 |
+
def render_tasks(user_id):
|
| 669 |
+
if not user_id:
|
| 670 |
+
return ""
|
| 671 |
+
|
| 672 |
+
tasks = get_user_tasks(user_id)
|
| 673 |
+
|
| 674 |
+
if not tasks:
|
| 675 |
+
return "<div style='text-align: center; padding: 40px; color: #9ca3af;'>No tasks yet</div>"
|
| 676 |
+
|
| 677 |
+
html = ""
|
| 678 |
+
for task in tasks:
|
| 679 |
+
task_id, model, prompt, status, progress, tokens, tps, created = task
|
| 680 |
|
| 681 |
+
status_class = f"status-{status}"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 682 |
|
| 683 |
+
html += f"""
|
| 684 |
+
<div class="task-card">
|
| 685 |
+
<div style="display: flex; justify-content: space-between; margin-bottom: 8px;">
|
| 686 |
+
<strong>Task: {task_id[:8]}...</strong>
|
| 687 |
+
<span class="{status_class}">β{status.upper()}</span>
|
| 688 |
+
</div>
|
| 689 |
+
<div><strong>Model:</strong> {model}</div>
|
| 690 |
+
<div><strong>Prompt:</strong> {prompt[:100]}{'...' if len(prompt) > 100 else ''}</div>
|
| 691 |
+
<div class="progress-bar">
|
| 692 |
+
<div class="progress-fill" style="width: {progress}%"></div>
|
| 693 |
+
</div>
|
| 694 |
+
<div style="font-size: 12px; color: #6b7280;">
|
| 695 |
+
Progress: {progress}% | Tokens: {tokens} | Speed: {tps:.1f} tok/s
|
| 696 |
+
</div>
|
| 697 |
+
</div>
|
| 698 |
+
"""
|
| 699 |
+
|
| 700 |
+
return html
|
| 701 |
+
|
| 702 |
+
# Get task result
|
| 703 |
+
def get_task_result(user_id, task_id_short):
|
| 704 |
+
if not user_id or not task_id_short:
|
| 705 |
+
return "β Invalid request"
|
| 706 |
+
|
| 707 |
+
with db_lock:
|
| 708 |
+
c = db_conn.cursor()
|
| 709 |
+
c.execute("""SELECT result, status FROM tasks
|
| 710 |
+
WHERE user_id=? AND id LIKE ?""",
|
| 711 |
+
(user_id, f"{task_id_short}%"))
|
| 712 |
+
result = c.fetchone()
|
| 713 |
|
| 714 |
+
if result:
|
| 715 |
+
if result[1] == 'completed':
|
| 716 |
+
return f"### β
Result\n\n{result[0]}"
|
| 717 |
+
elif result[1] == 'failed':
|
| 718 |
+
return f"### β Failed\n\n{result[0]}"
|
| 719 |
+
else:
|
| 720 |
+
return f"### β³ Status: {result[1]}"
|
| 721 |
+
return "β Task not found"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 722 |
|
| 723 |
+
# Event handlers
|
| 724 |
+
auth_btn.click(
|
| 725 |
+
handle_auth,
|
| 726 |
+
[auth_username, auth_password],
|
| 727 |
+
[user_id_state, auth_msg, auth_group, main_group]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 728 |
)
|
| 729 |
|
| 730 |
+
submit_btn.click(
|
| 731 |
+
submit_task,
|
| 732 |
+
[user_id_state, model_choice, prompt_input],
|
| 733 |
+
[task_msg, prompt_input]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 734 |
).then(
|
| 735 |
+
render_tasks,
|
| 736 |
+
[user_id_state],
|
| 737 |
+
[tasks_display]
|
| 738 |
)
|
| 739 |
|
| 740 |
+
refresh_btn.click(
|
| 741 |
+
render_tasks,
|
| 742 |
+
[user_id_state],
|
| 743 |
+
[tasks_display]
|
|
|
|
|
|
|
|
|
|
| 744 |
)
|
| 745 |
|
| 746 |
+
# Auto-refresh timer
|
| 747 |
+
def auto_refresh_tasks(user_id, enabled):
|
| 748 |
+
if enabled and user_id:
|
| 749 |
+
return render_tasks(user_id)
|
| 750 |
+
return gr.update()
|
| 751 |
+
|
| 752 |
+
# Poll every 3 seconds when auto-refresh enabled
|
| 753 |
+
demo.load(
|
| 754 |
+
lambda: None,
|
| 755 |
+
None,
|
| 756 |
+
None,
|
| 757 |
+
every=3
|
| 758 |
)
|
| 759 |
|
| 760 |
+
# Update user display on load
|
| 761 |
+
def update_user_display(user_id):
|
| 762 |
+
if user_id:
|
| 763 |
+
with db_lock:
|
| 764 |
+
c = db_conn.cursor()
|
| 765 |
+
c.execute("SELECT username FROM users WHERE id=?", (user_id,))
|
| 766 |
+
result = c.fetchone()
|
| 767 |
+
if result:
|
| 768 |
+
active = get_user_active_tasks(user_id)
|
| 769 |
+
return f"**User:** {result[0]} | **Active:** {active}/5"
|
| 770 |
+
return ""
|
| 771 |
+
|
| 772 |
+
# Periodic refresh
|
| 773 |
+
refresh_timer = gr.Timer(3)
|
| 774 |
+
|
| 775 |
+
@refresh_timer.tick
|
| 776 |
+
def timer_refresh(user_id, auto_enabled):
|
| 777 |
+
if auto_enabled and user_id:
|
| 778 |
+
return render_tasks(user_id), update_user_display(user_id)
|
| 779 |
+
return gr.update(), gr.update()
|
| 780 |
+
|
| 781 |
+
refresh_timer.tick(
|
| 782 |
+
timer_refresh,
|
| 783 |
+
[user_id_state, auto_refresh],
|
| 784 |
+
[tasks_display, user_display]
|
| 785 |
+
)
|
| 786 |
|
| 787 |
+
# View full result (expandable)
|
| 788 |
+
with gr.Accordion("π View Task Result", open=False):
|
| 789 |
+
result_task_id = gr.Textbox(
|
| 790 |
+
label="Task ID (first 8 chars)",
|
| 791 |
+
placeholder="e.g., 3f7a9b2c"
|
| 792 |
+
)
|
| 793 |
+
view_result_btn = gr.Button("View Result", variant="primary")
|
| 794 |
+
result_display = gr.Markdown("")
|
| 795 |
|
| 796 |
+
view_result_btn.click(
|
| 797 |
+
get_task_result,
|
| 798 |
+
[user_id_state, result_task_id],
|
| 799 |
+
[result_display]
|
| 800 |
)
|
| 801 |
|
| 802 |
+
# Initial load
|
| 803 |
+
def on_auth_success(user_id):
|
| 804 |
+
if user_id:
|
| 805 |
+
return render_tasks(user_id), update_user_display(user_id)
|
| 806 |
+
return "", ""
|
| 807 |
+
|
| 808 |
+
user_id_state.change(
|
| 809 |
+
on_auth_success,
|
| 810 |
+
[user_id_state],
|
| 811 |
+
[tasks_display, user_display]
|
| 812 |
)
|
| 813 |
|
|
|
|
| 814 |
if __name__ == "__main__":
|
| 815 |
+
print("\n" + "="*80)
|
| 816 |
+
print("π SAM BACKGROUND PROCESSOR".center(80))
|
| 817 |
+
print("="*80)
|
| 818 |
+
print(f"β
2 worker threads active")
|
| 819 |
+
print(f"β
Max 5 tasks per user")
|
| 820 |
+
print(f"β
Background processing enabled")
|
| 821 |
+
print(f"β
Database: sam_tasks.db")
|
| 822 |
+
print("="*80 + "\n")
|
| 823 |
+
|
| 824 |
+
demo.queue(max_size=50)
|
| 825 |
demo.launch(
|
| 826 |
server_name="0.0.0.0",
|
| 827 |
server_port=7860,
|