Update app.py

app.py

@@ -1,142 +1,256 @@
+import os
+import torch
 import gradio as gr
-from huggingface_hub import InferenceClient, HfHubHTTPError, InferenceTimeoutError
-import httpx # 确保这个库在requirements.txt中
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from typing import List, Tuple, Dict
 
-"""
-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
-"""
+# 如果需要使用Hugging Face访问令牌，取消下面两行的注释并设置环境变量
+# from huggingface_hub import login
+# login(token=os.getenv("HUGGINGFACE_TOKEN"))
 
-# 定义多个模型及其对应的InferenceClient
-# !!重要!! 如果需要，请替换 'hf_YOUR_TOKEN_HERE' 为您的Hugging Face API Token。
-# API Token可以在 https://huggingface.co/settings/tokens 生成。
-# 使用API Token可以帮助解决一些访问限制问题，特别是对于热门模型。
-# 同时，请确保您已经在Hugging Face网站上同意了Mistral 7B Instruct v0.2等模型的条款。
-MODEL_CLIENTS = {
-    "Zephyr 7B Beta": InferenceClient("HuggingFaceH4/zephyr-7b-beta"), # 默认使用API Token
-    "Mistral 7B Instruct v0.2": InferenceClient("mistralai/Mistral-7B-Instruct-v0.2"), # 默认使用API Token
-    # 如果您需要使用API Token，可以这样写：
-    # "Zephyr 7B Beta": InferenceClient("HuggingFaceH4/zephyr-7b-beta", token="hf_YOUR_TOKEN_HERE"),
-    # "Mistral 7B Instruct v0.2": InferenceClient("mistralai/Mistral-7B-Instruct-v0.2", token="hf_YOUR_TOKEN_HERE"),
-    # 更多模型示例（请根据需要添加或删除）：
-    # "Llama 2 7B Chat": InferenceClient("meta-llama/Llama-2-7b-chat-hf", token="hf_YOUR_TOKEN_HERE"), # Llama 2通常需要访问权限和API Token
-    # "OpenHermes 2.5 Mistral 7B": InferenceClient("teknium/OpenHermes-2.5-Mistral-7B"),
+# 模型配置 - 可根据需要添加更多模型
+MODELS = {
+    "Llama 2 7B Chat": {
+        "model_id": "meta-llama/Llama-2-7b-chat-hf",
+        "kwargs": {"torch_dtype": torch.float16}
+    },
+    "Mistral 7B Instruct": {
+        "model_id": "mistralai/Mistral-7B-Instruct-v0.2",
+        "kwargs": {"torch_dtype": torch.float16}
+    },
+    "Zephyr 7B Beta": {
+        "model_id": "HuggingFaceH4/zephyr-7b-beta",
+        "kwargs": {"torch_dtype": torch.float16}
+    }
 }
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    # 新增参数，用于选择模型
-    selected_model_name,
-):
-    # 根据选择的模型名称获取对应的client
-    client = MODEL_CLIENTS.get(selected_model_name)
-    if not client:
-        # 如果模型名称无效，直接返回错误信息
-        yield "错误：未找到选定的模型客户端。请检查模型名称是否正确或已添加到列表中。"
-        return
-
-    messages = [{"role": "system", "content": system_message}]
-
-    # 构建完整的对话历史
-    for val in history:
-        if val[0]:  # 用户消息
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:  # 助手消息
-            messages.append({"role": "assistant", "content": val[1]})
+# 加载模型和分词器
+def load_model(model_name):
+    model_config = MODELS[model_name]
+    tokenizer = AutoTokenizer.from_pretrained(model_config["model_id"])
+    
+    # 检查模型是否需要特殊处理
+    if "Llama-2" in model_name:
+        model_config["kwargs"]["trust_remote_code"] = True
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        model_config["model_id"],
+        **model_config["kwargs"]
+    )
+    
+    # 将模型移动到可用设备
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(device)
+    
+    return model, tokenizer, device
 
-    messages.append({"role": "user", "content": message}) # 添加当前用户消息
-
-    response = ""
-    try:
-        # 使用选定的client进行推理
-        # client.chat_completion() 默认是一个生成器，用于流式传输
-        for message_chunk in client.chat_completion(
-            messages,
-            max_tokens=max_tokens,
-            stream=True, # 启用流式传输
-            temperature=temperature,
-            top_p=top_p,
-        ):
-            # 确保 chunk 和 content 存在，以防API响应格式异常
-            if message_chunk.choices and message_chunk.choices[0].delta and message_chunk.choices[0].delta.content is not None:
-                token = message_chunk.choices[0].delta.content
-                response += token
-                yield response # 逐步返回生成的文本
-            else:
-                # 可能是流的末尾，或者是一个空的内容块
-                pass
+# 初始化模型
+loaded_models = {}
+for model_name in MODELS:
+    loaded_models[model_name] = load_model(model_name)
 
-    # 错误处理：捕获可能出现的各种异常
-    except HfHubHTTPError as e:
-        error_message = ""
-        if e.response.status_code == 402:
-            error_message = "抱歉，此模型服务需要付费访问或您的Hugging Face账户额度已用尽。请检查您的Hugging Face账户设置或Space日志。"
-        elif e.response.status_code == 429:
-            error_message = "抱歉，请求过于频繁，触发了速率限制。请稍后再试，或考虑使用API Token提升额度。"
-        elif e.response.status_code == 401 or e.response.status_code == 403:
-            error_message = "抱歉，模型访问权限不足或API Token无效/缺失。请确保您已在Hugging Face上登录，接受模型条款，并正确配置API Token。"
-        elif e.response.status_code == 503:
-            error_message = "模型服务当前不可用，可能正在加载或维护中。请稍后再试。"
-        else:
-            error_message = f"模型服务出现HTTP错误 ({e.response.status_code})：{e.response.text}。请检查Hugging Face Space日志。"
-        print(f"HfHubHTTPError: {e}") # 打印到控制台以供调试
-        yield error_message # 将错误信息显示给用户
+# 构建对话提示词（针对Llama 2等需要特定格式的模型）
+def build_prompt(message, history, system_prompt):
+    prompt = f"[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n"
+    
+    # 添加对话历史
+    for user_msg, assistant_msg in history:
+        prompt += f"{user_msg} [/INST] {assistant_msg} [INST] "
+    
+    # 添加当前用户消息
+    prompt += f"{message} [/INST]"
+    
+    return prompt
 
-    except InferenceTimeoutError as e:
-        error_message = "模型响应超时，可能是请求过于复杂或服务器繁忙。请尝试减少'Max new tokens'或稍后再试。"
-        print(f"InferenceTimeoutError: {e}")
-        yield error_message
+# 模型推理函数
+def generate_response(
+    message: str,
+    history: List[Tuple[str, str]],
+    system_prompt: str,
+    model_name: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    top_k: int
+):
+    # 获取模型、分词器和设备
+    model, tokenizer, device = loaded_models[model_name]
+    
+    # 构建完整提示词
+    full_prompt = build_prompt(message, history, system_prompt)
+    
+    # 编码输入
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
+    
+    # 生成参数
+    generate_kwargs = {
+        "max_new_tokens": max_new_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "top_k": top_k,
+        "do_sample": True,
+        "eos_token_id": tokenizer.eos_token_id,
+        "pad_token_id": tokenizer.pad_token_id
+    }
+    
+    # 生成响应
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            **generate_kwargs
+        )
+    
+    # 解码输出
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    
+    # 提取模型生成的部分（去除提示词）
+    response = response[len(full_prompt):].strip()
+    
+    return response
 
-    except httpx.HTTPStatusError as e:
-        # 这是httpx库抛出的HTTP错误，可能发生在HfHubHTTPError之外
-        error_message = f"与Hugging Face服务通信时发生HTTP错误 ({e.response.status_code})：{e.response.text}。"
-        print(f"HTTPStatusError: {e}")
-        yield error_message
+# 处理用户输入并生成回复
+def process_chat(
+    message: str,
+    history: List[Tuple[str, str]],
+    system_prompt: str,
+    model_name: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    top_k: int
+):
+    # 生成响应
+    response = generate_response(
+        message, history, system_prompt, model_name,
+        max_new_tokens, temperature, top_p, top_k
+    )
+    
+    # 更新对话历史
+    history.append((message, response))
+    return history, history
 
+# 语音转文字功能（使用Whisper模型）
+asr = None
+if torch.cuda.is_available() or torch.backends.mps.is_available():
+    try:
+        from transformers import WhisperProcessor, WhisperForConditionalGeneration
+        processor = WhisperProcessor.from_pretrained("openai/whisper-base")
+        asr_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base").to("cuda" if torch.cuda.is_available() else "cpu")
+        asr = {
+            "processor": processor,
+            "model": asr_model
+        }
     except Exception as e:
-        # 捕获所有其他未预期的错误
-        error_message = f"发生未知错误：{type(e).__name__} - {e}。请查看Hugging Face Space日志了解更多详情。"
-        print(f"General Error: {e}")
-        yield error_message
-
+        print(f"语音识别模型加载失败: {e}")
+        asr = None
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs:
-https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="你是一个友好的AI助手，尽力提供帮助。", label="系统消息 (System message)"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="最大生成token数 (Max new tokens)"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="随机性 (Temperature)"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (核心采样)",
-        ),
-        # 新增一个Dropdown用于选择模型
-        gr.Dropdown(
-            list(MODEL_CLIENTS.keys()), # 选项为MODEL_CLIENTS的键（模型名称）
-            value=list(MODEL_CLIENTS.keys())[0], # 默认选中第一个模型
-            label="选择语言模型 (Select Model)",
-            interactive=True, # 允许用户更改
-        ),
-    ],
-    title="多模型AI聊天助手", # 给界面添加一个标题
-    description="选择一个语言模型，开始与AI对话。您可以调整参数或切换模型进行比较。", # 添加描述
-    submit_btn="发送",
-    stop_btn="停止",
-    clear_btn="清空对话",
-)
+def transcribe(audio):
+    if asr is None:
+        return "语音识别模型未加载"
+    
+    processor, model = asr["processor"], asr["model"]
+    input_features = processor(audio, return_tensors="pt").input_features.to(model.device)
+    predicted_ids = model.generate(input_features)
+    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
+    return transcription
 
+# 构建Gradio界面
+with gr.Blocks(title="语言模型对话助手") as demo:
+    gr.Markdown("## 基于Transformer的语言模型对话应用")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            # 输入区域
+            message_input = gr.Textbox(
+                label="输入消息", 
+                placeholder="请输入您想与AI对话的内容..."
+            )
+            
+            # 系统提示词
+            system_prompt = gr.Textbox(
+                label="系统提示词",
+                value="你是一个 helpful、知识渊博的AI助手。",
+                placeholder="设置AI的角色和行为准则..."
+            )
+            
+            # 模型选择
+            model_choice = gr.Dropdown(
+                choices=list(MODELS.keys()),
+                value=list(MODELS.keys())[0],
+                label="选择语言模型"
+            )
+            
+            # 生成参数
+            with gr.Accordion("高级生成参数", open=False):
+                max_new_tokens = gr.Slider(
+                    minimum=1, maximum=2048, value=512, step=1,
+                    label="最大生成Token数"
+                )
+                temperature = gr.Slider(
+                    minimum=0.1, maximum=2.0, value=0.7, step=0.1,
+                    label="温度（随机性）"
+                )
+                top_p = gr.Slider(
+                    minimum=0.1, maximum=1.0, value=0.9, step=0.05,
+                    label="Top-p（核采样）"
+                )
+                top_k = gr.Slider(
+                    minimum=1, maximum=100, value=50, step=1,
+                    label="Top-k（采样数）"
+                )
+            
+            # 语音输入
+            use_voice = gr.Checkbox(label="使用语音输入")
+            audio_input = gr.Audio(
+                type="filepath", 
+                label="语音输入（录制或上传音频）"
+            )
+            
+            # 按钮
+            send_btn = gr.Button("发送消息", variant="primary")
+            clear_btn = gr.Button("清空对话")
+        
+        with gr.Column(scale=2):
+            # 对话历史
+            chat_history = gr.Chatbot(
+                label="对话历史",
+                show_label=True
+            )
+    
+    # 语音输入处理
+    def handle_voice(audio, use_voice):
+        if use_voice and audio:
+            return transcribe(audio)
+        return ""
+    
+    audio_input.change(
+        fn=handle_voice,
+        inputs=[audio_input, use_voice],
+        outputs=message_input
+    )
+    
+    # 发送消息处理
+    send_btn.click(
+        fn=process_chat,
+        inputs=[
+            message_input, chat_history, system_prompt, model_choice,
+            max_new_tokens, temperature, top_p, top_k
+        ],
+        outputs=[chat_history, chat_history],
+        show_progress=True
+    )
+    
+    # 清空对话
+    clear_btn.click(
+        fn=lambda: None,
+        inputs=None,
+        outputs=chat_history
+    )
 
+# 启动应用
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(
+        server_name="0.0.0.0", 
+        server_port=7860, 
+        share=True
+    )