Update app.py

app.py

@@ -1,50 +1,117 @@
 import gradio as gr
 import requests
 import time
+import os
+from transformers import pipeline, set_seed
 
-def generate_protein(instruction, max_length=100):
-    """生成蛋白质序列"""
+# 设置随机种子
+set_seed(42)
+
+# 从环境变量获取 token（现在应该可以正常工作了）
+HF_TOKEN = os.environ.get('ProtTex')
+print(f"Token 状态: {'已设置' if HF_TOKEN else '未设置'}")
+
+# 初始化模型管道
+try:
+    protein_generator = pipeline(
+        "text-generation", 
+        model="mzcwd/ProtTeX",
+        device=-1,  # 使用CPU
+        torch_dtype="auto"
+    )
+    MODEL_LOADED = True
+    print("✅ 本地模型加载成功")
+except Exception as e:
+    print(f"❌ 本地模型加载失败: {e}")
+    MODEL_LOADED = False
+    protein_generator = None
+
+def generate_with_local_model(instruction, max_length=100):
+    """使用本地加载的模型生成蛋白质序列"""
+    if not MODEL_LOADED or protein_generator is None:
+        return "❌ 模型未正确加载，请检查控制台日志"
     
-    # 显示加载状态
-    yield "正在生成蛋白质序列，请稍候..."
+    try:
+        result = protein_generator(
+            instruction,
+            max_length=max_length,
+            num_return_sequences=1,
+            temperature=0.8,
+            do_sample=True,
+            pad_token_id=50256
+        )
+        
+        if result and len(result) > 0:
+            return f"✅ 生成成功:\n\n{result[0]['generated_text']}"
+        else:
+            return "❌ 生成失败，未获得有效结果"
+            
+    except Exception as e:
+        return f"❌ 生成过程中出现错误: {str(e)}"
+
+def generate_with_api(instruction, max_length=100):
+    """使用HuggingFace API生成蛋白质序列"""
+    if not HF_TOKEN:
+        return "❌ 未设置 HuggingFace Token，无法使用 API 功能"
     
     try:
         API_URL = "https://api-inference.huggingface.co/models/mzcwd/ProtTeX"
         
+        headers = {
+            "Authorization": f"Bearer {HF_TOKEN}",
+            "Content-Type": "application/json"
+        }
+        
         payload = {
             "inputs": instruction,
             "parameters": {
-                "max_length": max_length,
+                "max_new_tokens": max_length,
                 "temperature": 0.8,
                 "do_sample": True,
                 "return_full_text": False
-            },
-            "options": {
-                "wait_for_model": True
             }
         }
         
-        # 发送请求到 HuggingFace 推理 API
-        response = requests.post(API_URL, json=payload)
+        print("🔄 正在调用 API...")
+        response = requests.post(API_URL, headers=headers, json=payload, timeout=60)
         
         if response.status_code == 200:
             result = response.json()
             if isinstance(result, list) and len(result) > 0:
-                generated_text = result[0].get('generated_text', '')
-                yield generated_text if generated_text else "模型已响应但未生成内容"
-            else:
-                yield f"响应格式异常: {result}"
-        elif response.status_code == 503:
-            # 模型正在加载，等待后重试
-            yield "模型正在启动中，请等待约30秒后重试..."
+                generated_text = result[0].get('generated_text', '生成成功但无内容返回')
+                return f"✅ API 生成成功:\n\n{generated_text}"
+            return f"✅ API 响应: {str(result)}"
         else:
-            error_msg = f"请求失败 (状态码: {response.status_code})"
-            if response.status_code == 401:
-                error_msg += "\n可能需要设置访问令牌"
-            yield error_msg
+            error_msg = f"❌ API请求失败 (状态码: {response.status_code})"
+            if response.text:
+                error_msg += f"\n详细信息: {response.text}"
+            return error_msg
             
+    except requests.exceptions.Timeout:
+        return "⏰ 请求超时，请稍后重试"
     except Exception as e:
-        yield f"发生错误: {str(e)}"
+        return f"❌ API调用错误: {str(e)}"
+
+def generate_protein(instruction, max_length=100):
+    """主生成函数"""
+    yield "🔄 正在处理您的请求..."
+    time.sleep(0.5)
+    
+    try:
+        # 显示当前模式
+        mode = "本地模型" if MODEL_LOADED else "API调用"
+        yield f"🔄 使用 {mode} 处理中..."
+        
+        # 优先使用本地模型，失败时使用API
+        if MODEL_LOADED:
+            result = generate_with_local_model(instruction, max_length)
+        else:
+            result = generate_with_api(instruction, max_length)
+        
+        yield result
+        
+    except Exception as e:
+        yield f"❌ 生成过程中出现未预期错误: {str(e)}"
 
 # 创建界面
 with gr.Blocks(
@@ -52,7 +119,8 @@ with gr.Blocks(
     theme=gr.themes.Soft(),
     css="""
     .gradio-container {
-        max-width: 800px !important;
+        max-width: 900px;
+        margin: auto;
     }
     """
 ) as demo:
@@ -60,16 +128,18 @@ with gr.Blocks(
     gr.Markdown("""
     # 🧬 ProtTeX 蛋白质生成器
     
-    使用自然语言指令生成蛋白质序列。输入您想要的蛋白质特性描述，AI将生成相应的蛋白质序列。
+    **使用自然语言指令生成蛋白质序列**
+    
+    *当前模式: {'🧠 本地模型' if MODEL_LOADED else '🌐 API 调用'} | Token: {'✅ 已设置' if HF_TOKEN else '❌ 未设置'}*
     """)
     
     with gr.Row():
         with gr.Column(scale=1):
             instruction = gr.Textbox(
-                label="蛋白质生成指令",
+                label="🧪 蛋白质生成指令",
                 placeholder="例如：Generate a protein with alpha-helical structure for membrane binding",
                 lines=3,
-                max_lines=5
+                value="Generate a hydrophobic transmembrane protein sequence"
             )
             
             max_length = gr.Slider(
@@ -77,7 +147,7 @@ with gr.Blocks(
                 maximum=300,
                 value=150,
                 step=10,
-                label="序列最大长度"
+                label="📏 序列最大长度"
             )
             
             generate_btn = gr.Button(
@@ -88,14 +158,14 @@ with gr.Blocks(
             
         with gr.Column(scale=1):
             output = gr.Textbox(
-                label="生成的蛋白质序列",
-                lines=6,
+                label="🧬 生成的蛋白质序列",
+                lines=8,
                 show_copy_button=True,
-                interactive=False
+                elem_id="output-box"
             )
     
     # 示例部分
-    gr.Markdown("### 💡 示例指令")
+    gr.Markdown("### 💡 示例指令（点击试用）")
     
     examples = gr.Examples(
         examples=[
@@ -105,10 +175,7 @@ with gr.Blocks(
             ["Generate a stable protein for high temperature environments"]
         ],
         inputs=[instruction],
-        outputs=[output],
-        fn=generate_protein,
-        cache_examples=False,
-        label="点击示例快速尝试"
+        label="点击下面的示例快速尝试"
     )
     
     # 连接按钮事件
@@ -118,14 +185,20 @@ with gr.Blocks(
         outputs=[output]
     )
     
-    # 页脚信息
-    gr.Markdown("""
+    # 状态信息
+    gr.Markdown(f"""
     ---
+    **系统状态**:
+    - 运行模式: {'🧠 本地模型' if MODEL_LOADED else '🌐 API 调用'}
+    - Token 状态: {'✅ 已设置' if HF_TOKEN else '❌ 未设置'}
+    - 硬件: CPU Basic
+    
     *基于 [mzcwd/ProtTeX](https://huggingface.co/mzcwd/ProtTeX) 模型*
     """)
 
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
-        share=False
+        share=False,
+        show_error=True
     )