Update app.py

app.py

@@ -1,15 +1,11 @@
 import gradio as gr
-from huggingface_hub import InferenceClient, hf_hub_download
-from transformers import AutoTokenizer
-import os
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# --- 设置模型 ID ---
+# --- 直接加载模型和分词器 ---
 model_id = "xingyu1996/tiger-gpt2"
-client = InferenceClient(model_id)
-
-# --- 关键变化: 直接加载与训练时相同的分词器 ---
-tokenizer = AutoTokenizer.from_pretrained("gpt2")
+tokenizer = AutoTokenizer.from_pretrained("gpt2") # 使用原始的 GPT-2 分词器
+model = AutoModelForCausalLM.from_pretrained(model_id)
 
 def respond(
     message,
@@ -18,42 +14,30 @@ def respond(
     temperature,
     top_p,
 ):
-    prompt = message
-    response_ids = []
-    response_text = ""
-
-    # --- 参数准备 ---
-    generation_args = {
-        "max_new_tokens": max_tokens,
-        "stream": True,
-        "details": True,  # 让 API 返回 token ID (重要变化)
+    # 将输入文本转换为 token ID
+    input_ids = tokenizer.encode(message, return_tensors="pt")
+    
+    # 准备生成参数
+    gen_kwargs = {
+        "max_length": input_ids.shape[1] + max_tokens,
+        "do_sample": True if temperature > 0 else False,
     }
-
-    if temperature is not None and temperature > 0:
-        generation_args["temperature"] = temperature
-    if top_p is not None and top_p < 1.0:
-        generation_args["top_p"] = top_p
-
-    try:
-        # --- 调用 API, 获取 token IDs ---
-        for output in client.text_generation(prompt, **generation_args):
-            if hasattr(output, 'token'):  # 流式生成时的输出
-                # 这里 output.token 是一个字典，包含 id 和 text
-                token_id = output.token.id
-                response_ids.append(token_id)
-                
-                # 使用我们自己的 tokenizer 来解码
-                current_text = tokenizer.decode(response_ids, skip_special_tokens=True)
-                response_text = current_text
-                
-                yield response_text
-            elif hasattr(output, 'generated_text'):  # 非流式生成时的最终输出
-                # 如果直接返回了完整文本，就用它
-                response_text = output.generated_text
-                yield response_text
-    except Exception as e:
-        print(f"推理时发生错误: {type(e).__name__} - {e}")
-        yield f"抱歉，推理时遇到错误: {type(e).__name__} - {str(e)}"
+    
+    if temperature > 0:
+        gen_kwargs["temperature"] = temperature
+    if top_p < 1.0:
+        gen_kwargs["top_p"] = top_p
+        
+    # 生成文本
+    with torch.no_grad():
+        output_ids = model.generate(input_ids, **gen_kwargs)
+        
+    # 只保留新生成的部分
+    new_tokens = output_ids[0, input_ids.shape[1]:]
+    
+    # 解码生成的 token ID
+    response = tokenizer.decode(new_tokens, skip_special_tokens=True)
+    return response
 
 
 # 其他 Gradio 界面代码不变