Update app.py

app.py

@@ -20,12 +20,6 @@ import outlines
 from outlines.types import Regex  # 导入最新的 Regex 类型限制
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# 【新增】生成当前合法走法的正则表达式
-def generate_regex(board: chess.Board) -> str:
-    legal_moves_san = [board.san(m) for m in board.legal_moves]
-    pattern = "|".join(re.escape(san) for san in legal_moves_san)
-    return f" ?({pattern})"
-
 
 # ──────────────────────────────────────────────────────────────────────────────
 # Available models
@@ -41,47 +35,28 @@ MODEL_DESCRIPTIONS = {
 
 device = torch.device("cpu")
 
-# Lazy cache: {model_id: (tokenizer, model)}
-# _model_cache: dict = {}
+Lazy cache: {model_id: (tokenizer, model)}
+_model_cache: dict = {}
 
 @spaces.GPU(duration=30)
 def foo(bar):
     return bar
 
-_MODEL_CACHE = {}
-
 
 def load_model(model_key: str):
+    """Load (or retrieve from cache) tokenizer + model for the given key."""
     model_id = AVAILABLE_MODELS[model_key]
-    
-    # 只有第一次使用该模型时才会加载模型权重
-    if model_id not in _MODEL_CACHE:
-        print(f"🚀 [首次加载] 正在加载模型权重到内存: {model_id} …")
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        hf_model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
-        
-        # 保存包装后的 Outlines 模型对象
-        _MODEL_CACHE[model_id] = outlines.from_transformers(hf_model, tokenizer)
-        print(f"✓ {model_id} 已成功缓存，后续走棋将直接复用内存！")
-        
-    return _MODEL_CACHE[model_id]
-# _model_cache: dict = {}
-
-
-# def load_model(model_key: str):
-#     """Load (or retrieve from cache) tokenizer + model for the given key."""
-#     model_id = AVAILABLE_MODELS[model_key]
-#     if model_id not in _model_cache:
-#         print(f"Loading {model_id} …")
-#         tokenizer = GPT2Tokenizer.from_pretrained(model_id)
-#         tokenizer.pad_token = tokenizer.eos_token
-#         model = GPT2LMHeadModel.from_pretrained(model_id)
-#         model.to(device)
-#         model.eval()
-#         model.config.use_cache = True
-#         _model_cache[model_id] = (tokenizer, model)
-#         print(f"✓ {model_id} ready on {device}")
-#     return _model_cache[model_id]
+    if model_id not in _model_cache:
+        print(f"Loading {model_id} …")
+        tokenizer = GPT2Tokenizer.from_pretrained(model_id)
+        tokenizer.pad_token = tokenizer.eos_token
+        model = GPT2LMHeadModel.from_pretrained(model_id)
+        model.to(device)
+        model.eval()
+        model.config.use_cache = True
+        _model_cache[model_id] = (tokenizer, model)
+        print(f"✓ {model_id} ready on {device}")
+    return _model_cache[model_id]
 
 # ──────────────────────────────────────────────────────────────────────────────
 # Chess / model logic
@@ -130,48 +105,29 @@ def extract_move(text: str, board: chess.Board):
 
 @torch.no_grad()
 def get_model_move(board: chess.Board, model_key: str):
-    # 直接从全局缓存中获取已加载好的模型
-    model = load_model(model_key)
-    
-    prompt = board_to_prompt(board)
+    tokenizer, model = load_model(model_key)
+    prompt  = board_to_prompt(board)
     print("Current prompt: "+prompt)
 
-    regex_pattern = generate_regex(board)
-    
-    # 直接进行约束生成，不会触发模型重载
-    move_san = model(prompt, Regex(regex_pattern), do_sample=True, temperature=0.3, max_new_tokens=10)
-    move_san = str(move_san).strip()
-    
-    move = board.parse_san(move_san)
-    return move, True
-
-
-# @spaces.GPU(duration=10)
-# @torch.no_grad()
-# def get_model_move(board: chess.Board, model_key: str):
-#     tokenizer, model = load_model(model_key)
-#     prompt  = board_to_prompt(board)
-#     print("Current prompt: "+prompt)
-# 
-#     inputs  = tokenizer(prompt, return_tensors="pt").to(device)
-#     outputs = model.generate(
-#         inputs.input_ids,
-#         max_new_tokens=12,
-#         do_sample=True,
-#         temperature=0.3,
-#         top_k=40,
-#         top_p=0.9,
-#         repetition_penalty=1.1,
-#         pad_token_id=tokenizer.eos_token_id,
-#         eos_token_id=tokenizer.eos_token_id,
-#     )
-#     new_tokens = outputs[0][inputs.input_ids.shape[1]:]
-#     generated  = tokenizer.decode(new_tokens, skip_special_tokens=True)
-#     move       = extract_move(generated, board)
-#     if move:
-#         return move, True
-#     print(f"Wrong move: {generated}")
-#     return random.choice(list(board.legal_moves)), False
+    inputs  = tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = model.generate(
+        inputs.input_ids,
+        max_new_tokens=12,
+        do_sample=True,
+        temperature=0.3,
+        top_k=40,
+        top_p=0.9,
+        repetition_penalty=1.1,
+        pad_token_id=tokenizer.eos_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+    new_tokens = outputs[0][inputs.input_ids.shape[1]:]
+    generated  = tokenizer.decode(new_tokens, skip_special_tokens=True)
+    move       = extract_move(generated, board)
+    if move:
+        return move, True
+    print(f"Wrong move: {generated}")
+    return random.choice(list(board.legal_moves)), False