# modules/ai_model.py import torch import base64 import requests from io import BytesIO import os from huggingface_hub import login from PIL import Image from transformers import AutoProcessor, Gemma3nForConditionalGeneration from utils.logger import log from typing import Union, Tuple class AIModel: def __init__(self, model_name: str = "google/gemma-3n-e2b-it"): self.model_name = model_name self.model = None self.processor = None # 设置缓存目录 self._setup_cache_dirs() self._initialize_model() def _setup_cache_dirs(self): """设置缓存目录""" cache_dir = "/app/.cache/huggingface" os.makedirs(cache_dir, exist_ok=True) # 设置环境变量 os.environ["HF_HOME"] = cache_dir os.environ["TRANSFORMERS_CACHE"] = cache_dir os.environ["HF_DATASETS_CACHE"] = cache_dir log.info(f"设置缓存目录: {cache_dir}") def _authenticate_hf(self): assitant_token = os.getenv("Assitant_tocken") token_to_use = assitant_token cache_dir = "/app/.cache/huggingface" login(token=token_to_use, add_to_git_credential=False) log.info("✅ HuggingFace 认证成功") return token_to_use def _initialize_model(self): """初始化Gemma模型""" try: log.info(f"正在加载模型: {self.model_name}") token = self._authenticate_hf() if not token: log.error("❌ 无法获取有效token，模型加载失败") self.model = None self.processor = None return cache_dir = "/app/.cache/huggingface" self.model = Gemma3nForConditionalGeneration.from_pretrained( self.model_name, device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True, token=token, cache_dir=cache_dir ).eval() self.processor = AutoProcessor.from_pretrained( self.model_name, trust_remote_code=True, token=token, cache_dir=cache_dir ) log.info("✅ Gemma AI 模型初始化成功") except Exception as e: log.error(f"❌ Gemma AI 模型初始化失败: {e}", exc_info=True) self.model = None self.processor = None def is_available(self) -> bool: return self.model is not None and self.processor is not None def detect_input_type(self, input_data: str) -> str: if not isinstance(input_data, str): return "text" image_extensions = [".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp"] if (input_data.startswith(("http://", "https://")) and any(input_data.lower().endswith(ext) for ext in image_extensions)): return "image" elif any(input_data.endswith(ext) for ext in image_extensions): return "image" elif input_data.startswith("data:image/"): return "image" audio_extensions = [".wav", ".mp3", ".m4a", ".ogg", ".flac"] if (input_data.startswith(("http://", "https://")) and any(input_data.lower().endswith(ext) for ext in audio_extensions)): return "audio" elif any(input_data.endswith(ext) for ext in audio_extensions): return "audio" return "text" def format_input(self, input_type: str, raw_input: str) -> Tuple[str, Union[str, Image.Image, None]]: if input_type == "image": try: if raw_input.startswith("data:image/"): header, encoded = raw_input.split(",", 1) image_data = base64.b64decode(encoded) image = Image.open(BytesIO(image_data)).convert("RGB") elif raw_input.startswith(("http://", "https://")): response = requests.get(raw_input, timeout=10) response.raise_for_status() image = Image.open(BytesIO(response.content)).convert("RGB") else: image = Image.open(raw_input).convert("RGB") log.info("✅ 图片加载成功") return input_type, image, "请描述这张图片，并基于图片内容提供旅游建议。" except Exception as e: log.error(f"❌ 图片加载失败: {e}") return "text", None, f"图片加载失败，请检查路径或URL。" elif input_type == "audio": log.warning("⚠️ 音频处理功能暂未实现") return "text", None, "抱歉，音频输入功能正在开发中。请使用文字描述您的需求。" else: # text return input_type, None, raw_input def run_inference(self, input_type: str, formatted_input: Union[str, Image.Image], prompt: str,temperature: float = 0.5) -> str: try: # 截断过长的 prompt if len(prompt) > 500: prompt = prompt[:500] + "..." # 准备输入 (处理图片或文本) if input_type == "image" and isinstance(formatted_input, Image.Image): image_token = getattr(self.processor.tokenizer, 'image_token', '') if image_token not in prompt: prompt = f"{image_token}\n{prompt}" inputs = self.processor( text=prompt, images=formatted_input, return_tensors="pt" ).to(self.model.device, dtype=torch.bfloat16) else: inputs = self.processor( text=prompt, return_tensors="pt" ).to(self.model.device, dtype=torch.bfloat16) if hasattr(inputs, 'input_ids') and inputs.input_ids.shape[-1] > 512: log.warning(f"⚠️ 截断过长输入: {inputs.input_ids.shape[-1]} -> 512") inputs.input_ids = inputs.input_ids[:, :512] if hasattr(inputs, 'attention_mask'): inputs.attention_mask = inputs.attention_mask[:, :512] with torch.inference_mode(): generation_args = { "max_new_tokens": 512, "pad_token_id": self.processor.tokenizer.eos_token_id, "use_cache": True } # 如果 temperature 接近0，使用贪心解码 (用于分类等确定性任务) if temperature < 1e-6: log.info("▶️ 使用贪心解码 (do_sample=False) 以获得确定性输出。") generation_args["do_sample"] = False # 否则，使用采样解码 (用于创造性生成任务) else: log.info(f"▶️ 使用采样解码 (do_sample=True)，temperature={temperature}。") generation_args["do_sample"] = True generation_args["temperature"] = temperature generation_args["top_p"] = 0.9 # top_p 只在采样时有意义 # 使用构建好的参数字典来调用 generate outputs = self.model.generate( **inputs, **generation_args ) input_length = inputs.input_ids.shape[-1] generated_tokens = outputs[0][input_length:] decoded = self.processor.tokenizer.decode(generated_tokens, skip_special_tokens=True).strip() return decoded if decoded else "我理解了您的问题，请告诉我更多具体信息。" except RuntimeError as e: if "shape" in str(e): log.error(f"❌ Tensor形状错误: {e}") return "输入处理遇到问题，请尝试简化您的问题。" raise e except Exception as e: log.error(f"❌ 模型推理失败: {e}", exc_info=True) return "抱歉，处理您的请求时遇到技术问题。" def chat_completion(self, model: str, messages: list, **kwargs) -> str: if not self.is_available(): log.error("模型未就绪，无法执行 chat_completion") if kwargs.get("response_format", {}).get("type") == "json_object": return '{"error": "Model not available"}' return "抱歉，AI 模型当前不可用。" full_prompt = "\n".join([msg.get("content", "") for msg in messages]) temperature = kwargs.get("temperature", 0.7) if kwargs.get("response_format", {}).get("type") == "json_object": # 在 prompt 末尾添加指令，强制模型输出 JSON full_prompt += "\n\n请注意：你的回答必须是一个严格的、不含任何额外解释和代码块标记的 JSON 对象。" # 对于JSON生成任务，使用较低的 temperature 以获得更稳定、确定性的结构 temperature = 0.1 log.debug(f"▶️ 执行 chat_completion (适配器), temperature={temperature}, prompt='{full_prompt[:100]}...'") return self.run_inference( input_type="text", formatted_input=None, prompt=full_prompt, temperature=temperature # 将处理后的 temperature 传递下去 ) def _build_limited_prompt(self, processed_text: str, context: str = "") -> str: """构建长度受限的prompt - 新增辅助方法""" # 限制输入长度 if len(processed_text) > 200: processed_text = processed_text[:200] + "..." if context and len(context) > 300: context = context[:300] + "..." # 保持你原有的prompt结构 if context: return ( f"你是一个专业的旅游助手。请基于以下背景信息，用中文友好地回答用户的问题。\n\n" f"--- 背景信息 ---\n{context}\n\n" f"--- 用户问题 ---\n{processed_text}\n\n" f"请提供专业、实用的旅游建议：" ) else: return ( f"你是一个专业的旅游助手。请用中文友好地回答用户的问题。\n\n" f"用户问题：{processed_text}\n\n" f"请提供专业、实用的旅游建议：" ) def generate(self, user_input: str, context: str = "") -> str: """主要的生成方法 - 保持原有逻辑""" if not self.is_available(): return "抱歉，AI 模型当前不可用，请稍后再试。" try: # 1. 检测输入类型 input_type = self.detect_input_type(user_input) log.info(f"检测到输入类型: {input_type}") # 2. 格式化输入 input_type, formatted_data, processed_text = self.format_input(input_type, user_input) # 3. 构建prompt - 使用你的原有结构 prompt = self._build_limited_prompt(processed_text, context) # 4. 执行推理 if input_type == "image" and formatted_data is not None: return self.run_inference("image", formatted_data, prompt) else: return self.run_inference("text", processed_text, prompt) except Exception as e: log.error(f"❌ 生成回复时发生错误: {e}", exc_info=True) return "抱歉，我在思考时遇到了点麻烦，请稍后再试。"