Update app.py

app.py

@@ -6,9 +6,11 @@ from transformers import AutoTokenizer
 import torch.nn as nn
 import torch.nn.functional as F
 
-# ================ 第一步：重新定义模型结构 (必须与训练时完全一致) ================
-# 注意：这些类定义必须与你原始训练脚本中的完全相同
+"""
+For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
+"""
 
+# ================ 第一步：定义模型结构 ================
 class GELU(nn.Module):
     def __init__(self):
         super().__init__()
@@ -139,20 +141,64 @@ class GPTModel(nn.Module):
         logits = self.out_head(x)
         return logits
 
-# 用于生成的函数
-def generate_text_simple(model, idx, max_new_tokens, context_size):
+# ================ 第二步：定义文本生成函数 ================
+
+def generate_text_simple(model, idx, max_new_tokens, context_size, temperature=1.0, top_k=None):
+    """
+    使用 top_k 采样和温度缩放的文本生成函数
+    
+    参数:
+        model: 语言模型
+        idx: 输入序列的 token ID
+        max_new_tokens: 要生成的最大新 token 数量
+        context_size: 上下文窗口大小
+        temperature: 温度参数，控制采样的随机性（越高越随机）
+        top_k: 只考虑概率最高的 top_k 个 token，如果为 None 或 0 则考虑所有 token
+    
+    返回:
+        扩展后的 token ID 序列
+    """
     device = idx.device
-    current_device_type = str(device).split(':')[0]
     
     for _ in range(max_new_tokens):
+        # 获取当前上下文
         idx_cond = idx[:, -context_size:]
+        
         with torch.no_grad():
-            # 推理时不需要混合精度
+            # 获取模型预测的下一个 token 的 logits
             logits = model(idx_cond)
-        logits = logits[:, -1, :]
-        probas = torch.softmax(logits, dim=-1)
-        idx_next = torch.argmax(probas, dim=-1, keepdim=True)
+            # 只关心最后一个位置的预测
+            logits = logits[:, -1, :]
+            
+            # 应用温度缩放
+            if temperature > 0:
+                logits = logits / temperature
+            
+            # 应用 top_k 过滤
+            if top_k is not None and top_k > 0:
+                # 获取前 k 个最大值
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                # 设置阈值为第 k 个最大值
+                threshold = v[..., [-1]]
+                # 将阈值以下的值设为 -inf
+                logits = torch.where(logits < threshold, 
+                                    torch.full_like(logits, float('-inf')), 
+                                    logits)
+            
+            # 应用 softmax 转换为概率
+            probs = torch.softmax(logits, dim=-1)
+            
+            # 根据概率分布采样
+            if temperature > 0:
+                # 随机采样
+                idx_next = torch.multinomial(probs, num_samples=1)
+            else:
+                # 如果温度为 0，则取最大概率的 token（等同于 argmax）
+                idx_next = torch.argmax(probs, dim=-1, keepdim=True)
+            
+        # 将新生成的 token 添加到序列中
         idx = torch.cat((idx, idx_next), dim=1)
+    
     return idx
 
 def text_to_token_ids(text, tokenizer):
@@ -164,7 +210,7 @@ def token_ids_to_text(token_ids, tokenizer):
     flat = token_ids.squeeze(0)
     return tokenizer.decode(flat.tolist(), skip_special_tokens=True)
 
-# ================ 第二步：设置模型加载和推理 ================
+# ================ 第三步：设置模型加载和推理 ================
 
 # 模型 ID
 model_id = "xingyu1996/tiger-gpt2"
@@ -190,7 +236,6 @@ def load_model_from_hub():
         config = json.load(f)
     
     # 将 Hugging Face 格式的配置转换为我们的格式
-    # 注意：这里的映射需要根据实际情况调整
     my_config = {
         "vocab_size": config.get("vocab_size", 50257),
         "context_length": config.get("n_positions", 512),
@@ -204,7 +249,6 @@ def load_model_from_hub():
     # 创建模型
     model = GPTModel(my_config)
     
-    # 加载权重到模型
     # 检查状态字典中是否有 _orig_mod. 前缀
     if any(k.startswith('_orig_mod.') for k in state_dict.keys()):
         state_dict = {k.replace('_orig_mod.', ''): v for k, v in state_dict.items()}
@@ -229,9 +273,9 @@ model, config = load_model_from_hub()
 tokenizer = AutoTokenizer.from_pretrained("gpt2")
 print("模型和分词器加载完成！")
 
-# ================ 第三步：设置 Gradio 接口 ================
+# ================ 第四步：设置 Gradio 接口 ================
 
-def respond(message, history, max_tokens, temperature):
+def respond(message, history, max_tokens, temperature, top_k):
     input_ids = text_to_token_ids(message, tokenizer).to("cpu")  # Hugging Face Space 可能没有 GPU
     context_size = config["context_length"]
     
@@ -241,7 +285,9 @@ def respond(message, history, max_tokens, temperature):
             model=model,
             idx=input_ids,
             max_new_tokens=max_tokens,
-            context_size=context_size
+            context_size=context_size,
+            temperature=temperature,
+            top_k=top_k
         )
         
         # 解码生成的文本
@@ -263,10 +309,17 @@ demo = gr.ChatInterface(
     respond,
     additional_inputs=[
         gr.Slider(minimum=1, maximum=100, value=30, step=1, label="生成长度"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="温度"),
+        gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.1, label="温度 (0.0 表示无随机性)"),
+        gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Top-K (0 表示不限制)"),
     ],
     title=f"Tiger-GPT2 推理测试",
-    description="输入中文文本，模型将生成后续内容。此演示直接加载了原始模型权重，与本地推理行为一致。",
+    description="""输入中文文本，模型将生成后续内容。此演示直接加载了原始模型权重，与本地推理行为一致。
+    
+**参数说明**:
+- **生成长度**: 要生成的最大token数量
+- **温度**: 控制生成随机性，值越高越随机，值为0时始终选择最可能的词
+- **Top-K**: 只从概率最高的K个词中选择下一个词，设为0则考虑所有词
+    """,
 )
 
 if __name__ == "__main__":