模型推理与优化

class UIGenerator:
    """UI生成器 - 核心推理类"""
    
    def __init__(self, config_path: str = "config/model_config.yaml"):
        """初始化生成器"""
        with open(config_path, 'r', encoding='utf-8') as f:
            self.config = yaml.safe_load(f)
        
        self.model = None
        self.tokenizer = None
        self.lora_model = None
        
        # 规则回退模板
        self.fallback_templates = self._load_fallback_templates()

def load_model(self, model_path: str, lora_path: Optional[str] = None):
    """加载模型"""
    logger.info(f"加载模型: {model_path}")
    
    # 加载分词器
    self.tokenizer = AutoTokenizer.from_pretrained(model_path)
    
    # 加载基础模型
    self.model = AutoModelForSeq2SeqLM.from_pretrained(
        model_path,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto" if torch.cuda.is_available() else None
    )
    
    # 如果有LoRA权重，加载LoRA
    if lora_path and Path(lora_path).exists():
        logger.info(f"加载LoRA权重: {lora_path}")
        self.lora_model = PeftModel.from_pretrained(self.model, lora_path)
        self.model = self.lora_model
    
    # 设置为评估模式
    self.model.eval()
    
    logger.info("模型加载完成")

def load_model_optimized(self, model_path: str, lora_path: Optional[str] = None):
    """优化的模型加载"""
    # 使用低内存模式加载
    model_kwargs = {
        "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32,
        "low_cpu_mem_usage": True,  # 低CPU内存使用
        "device_map": "auto" if torch.cuda.is_available() else None
    }
    
    # 如果GPU内存不足，使用CPU卸载
    if torch.cuda.is_available():
        gpu_memory = torch.cuda.get_device_properties(0).total_memory
        if gpu_memory < 8 * 1024**3:  # 小于8GB
            model_kwargs["device_map"] = {"": "cpu"}  # 强制使用CPU
    
    self.model = AutoModelForSeq2SeqLM.from_pretrained(
        model_path,
        **model_kwargs
    )

class ModelCache:
    """模型缓存管理器"""
    
    def __init__(self, max_models: int = 3):
        self.cache = {}
        self.max_models = max_models
        self.access_order = []
    
    def get_model(self, model_key: str):
        """获取缓存的模型"""
        if model_key in self.cache:
            # 更新访问顺序
            self.access_order.remove(model_key)
            self.access_order.append(model_key)
            return self.cache[model_key]
        return None
    
    def cache_model(self, model_key: str, model):
        """缓存模型"""
        if len(self.cache) >= self.max_models:
            # 移除最久未使用的模型
            oldest_key = self.access_order.pop(0)
            del self.cache[oldest_key]
        
        self.cache[model_key] = model
        self.access_order.append(model_key)

def generate_with_model(self, prompt: str) -> str:
    """使用模型生成UI-DSL"""
    if self.model is None or self.tokenizer is None:
        raise ValueError("模型未加载")
    
    # 编码输入
    inputs = self.tokenizer(
        prompt,
        max_length=self.config["model"]["max_length"],
        padding=True,
        truncation=True,
        return_tensors="pt"
    )
    
    # 移动到GPU（如果可用）
    if torch.cuda.is_available():
        inputs = {k: v.cuda() for k, v in inputs.items()}
    
    # 生成
    with torch.no_grad():
        outputs = self.model.generate(
            **inputs,
            max_length=self.config["model"]["max_length"],
            temperature=self.config["model"]["temperature"],
            top_p=self.config["model"]["top_p"],
            top_k=self.config["model"]["top_k"],
            do_sample=True,
            num_return_sequences=1,
            pad_token_id=self.tokenizer.pad_token_id,
            eos_token_id=self.tokenizer.eos_token_id
        )
    
    # 解码输出
    generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return generated_text

# 生成参数配置
generation_config = {
    "max_length": 512,        # 最大生成长度
    "temperature": 0.7,       # 温度参数，控制随机性
    "top_p": 0.9,            # 核采样参数
    "top_k": 50,             # Top-K采样
    "do_sample": True,        # 是否使用采样
    "num_return_sequences": 1, # 返回序列数量
    "repetition_penalty": 1.1, # 重复惩罚
    "length_penalty": 1.0,    # 长度惩罚
}

# 参数调优建议
parameter_tuning = {
    "temperature": {
        "range": [0.1, 1.5],
        "recommended": 0.7,
        "description": "温度越高，生成越随机；温度越低，生成越确定"
    },
    "top_p": {
        "range": [0.1, 1.0],
        "recommended": 0.9,
        "description": "核采样参数，控制词汇选择的多样性"
    },
    "top_k": {
        "range": [10, 100],
        "recommended": 50,
        "description": "Top-K采样，限制候选词汇数量"
    }
}

def batch_generate(self, prompts: List[str], batch_size: int = 4) -> List[str]:
    """批处理推理"""
    results = []
    
    for i in range(0, len(prompts), batch_size):
        batch_prompts = prompts[i:i + batch_size]
        
        # 编码批次输入
        inputs = self.tokenizer(
            batch_prompts,
            max_length=self.config["model"]["max_length"],
            padding=True,
            truncation=True,
            return_tensors="pt"
        )
        
        # 移动到GPU
        if torch.cuda.is_available():
            inputs = {k: v.cuda() for k, v in inputs.items()}
        
        # 批量生成
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_length=self.config["model"]["max_length"],
                temperature=self.config["model"]["temperature"],
                do_sample=True,
                num_return_sequences=1,
                pad_token_id=self.tokenizer.pad_token_id
            )
        
        # 解码批次输出
        batch_results = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
        results.extend(batch_results)
    
    return results

class InferenceCache:
    """推理缓存"""
    
    def __init__(self, max_size: int = 1000):
        self.cache = {}
        self.max_size = max_size
        self.access_count = {}
    
    def get(self, prompt: str) -> Optional[str]:
        """获取缓存结果"""
        if prompt in self.cache:
            self.access_count[prompt] = self.access_count.get(prompt, 0) + 1
            return self.cache[prompt]
        return None
    
    def put(self, prompt: str, result: str):
        """缓存结果"""
        if len(self.cache) >= self.max_size:
            # 移除访问次数最少的缓存
            min_access_prompt = min(self.access_count.keys(), key=lambda k: self.access_count[k])
            del self.cache[min_access_prompt]
            del self.access_count[min_access_prompt]
        
        self.cache[prompt] = result
        self.access_count[prompt] = 1

def _load_fallback_templates(self) -> Dict[str, Dict]:
    """加载规则回退模板"""
    return {
        "home": {
            "pattern": r"(首页|主页|电商首页)",
            "template": {
                "page": {
                    "name": "home_page",
                    "theme": "obsidian-gold",
                    "layout": {"grid": 12, "gutter": 16, "padding": 16, "bg": "#0E0E0E"},
                    "sections": [
                        {"type": "topbar", "props": {"logo": "品牌", "actions": ["search", "bell"]}},
                        {"type": "tabs", "props": {"items": ["出货", "求货"], "active": 0}},
                        {"type": "card-list", "props": {"columns": 2, "card": {"type": "product-card"}}},
                        {"type": "tabbar", "props": {"items": ["home", "search", "plus", "message", "user"]}}
                    ]
                }
            }
        },
        "detail": {
            "pattern": r"(详情|商品详情|详情页)",
            "template": {
                "page": {
                    "name": "detail_page",
                    "theme": "obsidian-gold",
                    "layout": {"grid": 12, "gutter": 16, "padding": 16, "bg": "#0E0E0E"},
                    "sections": [
                        {"type": "topbar", "props": {"title": "商品详情", "actions": ["share", "favorite"]}},
                        {"type": "carousel", "props": {"images": 3}},
                        {"type": "price", "props": {"value": "￥1999", "original": "￥2999"}},
                        {"type": "seller", "props": {"name": "优质卖家", "trust": "KYC已认证"}},
                        {"type": "proof", "props": {"items": ["发票", "订单截图", "序列号"]}},
                        {"type": "cta", "props": {"buttons": ["联系卖家", "立即下单"]}}
                    ]
                }
            }
        },
        "search": {
            "pattern": r"(搜索|搜索页|查找)",
            "template": {
                "page": {
                    "name": "search_page",
                    "theme": "obsidian-gold",
                    "layout": {"grid": 12, "gutter": 16, "padding": 16, "bg": "#0E0E0E"},
                    "sections": [
                        {"type": "topbar", "props": {"search": True, "actions": ["filter"]}},
                        {"type": "filters", "props": {"items": ["价格", "品牌", "地区"]}},
                        {"type": "card-list", "props": {"columns": 1, "card": {"type": "product-card"}}}
                    ]
                }
            }
        },
        "profile": {
            "pattern": r"(个人|个人中心|我的|用户中心)",
            "template": {
                "page": {
                    "name": "profile_page",
                    "theme": "obsidian-gold",
                    "layout": {"grid": 12, "gutter": 16, "padding": 16, "bg": "#0E0E0E"},
                    "sections": [
                        {"type": "topbar", "props": {"title": "个人中心", "actions": ["settings"]}},
                        {"type": "user-info", "props": {"avatar": True, "name": "用户名", "level": "VIP"}},
                        {"type": "menu-list", "props": {"items": ["我的订单", "我的收藏", "设置", "帮助"]}}
                    ]
                }
            }
        },
        "publish": {
            "pattern": r"(发布|发布页|发布商品)",
            "template": {
                "page": {
                    "name": "publish_page",
                    "theme": "obsidian-gold",
                    "layout": {"grid": 12, "gutter": 16, "padding": 16, "bg": "#0E0E0E"},
                    "sections": [
                        {"type": "topbar", "props": {"title": "发布商品", "actions": ["save"]}},
                        {"type": "form", "props": {"fields": ["title", "price", "description", "images"]}},
                        {"type": "cta", "props": {"buttons": ["保存草稿", "立即发布"]}}
                    ]
                }
            }
        }
    }

def generate_with_rules(self, prompt: str) -> str:
    """使用规则生成UI-DSL（回退方案）"""
    logger.info("使用规则回退生成")
    
    # 检测页面类型
    page_type = None
    for ptype, config in self.fallback_templates.items():
        if re.search(config["pattern"], prompt):
            page_type = ptype
            break
    
    if page_type is None:
        # 默认使用首页模板
        page_type = "home"
    
    # 获取模板
    template = self.fallback_templates[page_type]["template"].copy()
    
    # 根据Prompt调整模板
    template = self._adjust_template_by_prompt(template, prompt)
    
    # 转换为JSON字符串
    return json.dumps(template, ensure_ascii=False, separators=(',', ':'))

def _adjust_template_by_prompt(self, template: Dict, prompt: str) -> Dict:
    """根据Prompt调整模板"""
    # 检测主题
    if "黑金" in prompt or "obsidian-gold" in prompt:
        template["page"]["theme"] = "obsidian-gold"
        template["page"]["layout"]["bg"] = "#0E0E0E"
    elif "白银" in prompt or "silver-white" in prompt:
        template["page"]["theme"] = "silver-white"
        template["page"]["layout"]["bg"] = "#FFFFFF"
    elif "简约" in prompt or "minimal" in prompt:
        template["page"]["theme"] = "minimal"
        template["page"]["layout"]["bg"] = "#FFFFFF"
    
    # 检测列数
    if "单列" in prompt or "一列" in prompt:
        for section in template["page"]["sections"]:
            if section["type"] == "card-list":
                section["props"]["columns"] = 1
    elif "两列" in prompt or "2列" in prompt:
        for section in template["page"]["sections"]:
            if section["type"] == "card-list":
                section["props"]["columns"] = 2
    
    # 检测搜索栏
    if "搜索" in prompt:
        for section in template["page"]["sections"]:
            if section["type"] == "topbar":
                if "search" not in section["props"].get("actions", []):
                    section["props"]["actions"].append("search")
    
    return template

def validate_generated_json(self, generated_text: str) -> bool:
    """验证生成的JSON格式"""
    try:
        json.loads(generated_text)
        return True
    except json.JSONDecodeError:
        return False

def extract_json_from_text(self, text: str) -> Optional[str]:
    """从生成文本中提取JSON"""
    # 尝试直接解析
    try:
        json.loads(text)
        return text
    except:
        pass
    
    # 尝试提取JSON部分
    json_patterns = [
        r'\{.*\}',  # 匹配大括号内容
        r'```json\s*(\{.*?\})\s*```',  # 匹配代码块中的JSON
        r'```\s*(\{.*?\})\s*```',  # 匹配代码块
    ]
    
    for pattern in json_patterns:
        matches = re.findall(pattern, text, re.DOTALL)
        for match in matches:
            try:
                json.loads(match)
                return match
            except:
                continue
    
    return None

def check_content_quality(self, dsl: Dict) -> Dict[str, bool]:
    """检查生成内容质量"""
    quality_checks = {
        "has_required_fields": False,
        "has_valid_theme": False,
        "has_valid_sections": False,
        "has_reasonable_structure": False
    }
    
    # 检查必需字段
    if "page" in dsl and all(key in dsl["page"] for key in ["name", "theme", "layout", "sections"]):
        quality_checks["has_required_fields"] = True
    
    # 检查主题有效性
    valid_themes = ["obsidian-gold", "silver-white", "minimal"]
    if dsl.get("page", {}).get("theme") in valid_themes:
        quality_checks["has_valid_theme"] = True
    
    # 检查区块有效性
    valid_components = {"topbar", "tabs", "card-list", "carousel", "price", "seller", "proof", "cta", "tabbar", "user-info", "menu-list", "form", "filters"}
    sections = dsl.get("page", {}).get("sections", [])
    if sections and all(section.get("type") in valid_components for section in sections):
        quality_checks["has_valid_sections"] = True
    
    # 检查结构合理性
    if len(sections) >= 2 and len(sections) <= 8:
        quality_checks["has_reasonable_structure"] = True
    
    return quality_checks

def quantize_model(self, model, quantization_config=None):
    """模型量化"""
    if quantization_config is None:
        quantization_config = {
            "load_in_8bit": True,  # 8位量化
            "llm_int8_threshold": 6.0,
            "llm_int8_has_fp16_weight": False
        }
    
    # 应用量化
    quantized_model = BitsAndBytesConfig(**quantization_config)
    
    return quantized_model

def load_quantized_model(self, model_path: str):
    """加载量化模型"""
    quantization_config = BitsAndBytesConfig(
        load_in_8bit=True,
        llm_int8_threshold=6.0
    )
    
    self.model = AutoModelForSeq2SeqLM.from_pretrained(
        model_path,
        quantization_config=quantization_config,
        device_map="auto"
    )

class InferenceOptimizer:
    """推理优化器"""
    
    def __init__(self):
        self.optimization_config = {
            "use_cache": True,
            "use_kv_cache": True,
            "use_attention_cache": True,
            "batch_size": 4
        }
    
    def optimize_model(self, model):
        """优化模型推理"""
        # 启用KV缓存
        if hasattr(model, "use_cache"):
            model.use_cache = True
        
        # 编译模型（PyTorch 2.0+）
        if hasattr(torch, "compile"):
            model = torch.compile(model)
        
        return model
    
    def optimize_generation_config(self, generation_config):
        """优化生成配置"""
        optimized_config = generation_config.copy()
        
        # 启用缓存
        optimized_config["use_cache"] = True
        
        # 优化采样参数
        optimized_config["do_sample"] = True
        optimized_config["temperature"] = 0.7
        optimized_config["top_p"] = 0.9
        
        return optimized_config

import asyncio
from concurrent.futures import ThreadPoolExecutor

class AsyncInference:
    """异步推理"""
    
    def __init__(self, max_workers: int = 4):
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        self.model_lock = asyncio.Lock()
    
    async def generate_async(self, prompt: str) -> str:
        """异步生成"""
        loop = asyncio.get_event_loop()
        
        # 在线程池中执行推理
        result = await loop.run_in_executor(
            self.executor,
            self._generate_sync,
            prompt
        )
        
        return result
    
    def _generate_sync(self, prompt: str) -> str:
        """同步生成（在线程池中执行）"""
        # 实际的推理逻辑
        return self.generate_with_model(prompt)
    
    async def batch_generate_async(self, prompts: List[str]) -> List[str]:
        """异步批处理生成"""
        tasks = [self.generate_async(prompt) for prompt in prompts]
        results = await asyncio.gather(*tasks)
        return results

def generate(self, prompt: str, use_model: bool = True) -> Dict[str, Any]:
    """生成UI-DSL（带错误处理）"""
    try:
        if use_model and self.model is not None:
            # 使用模型生成
            generated_text = self.generate_with_model(prompt)
        else:
            # 使用规则生成
            generated_text = self.generate_with_rules(prompt)
        
        # 尝试解析JSON
        try:
            dsl = json.loads(generated_text)
            return {
                "success": True,
                "dsl": dsl,
                "raw_output": generated_text,
                "method": "model" if use_model and self.model is not None else "rules"
            }
        except json.JSONDecodeError:
            # JSON解析失败，使用规则回退
            logger.warning("模型输出JSON解析失败，使用规则回退")
            generated_text = self.generate_with_rules(prompt)
            dsl = json.loads(generated_text)
            return {
                "success": True,
                "dsl": dsl,
                "raw_output": generated_text,
                "method": "rules_fallback"
            }
    
    except Exception as e:
        logger.error(f"生成失败: {e}")
        return {
            "success": False,
            "error": str(e),
            "dsl": None,
            "raw_output": None,
            "method": "error"
        }

def safe_load_model(self, model_path: str, lora_path: Optional[str] = None) -> bool:
    """安全的模型加载"""
    try:
        self.load_model(model_path, lora_path)
        return True
    except Exception as e:
        logger.error(f"模型加载失败: {e}")
        
        # 尝试降级加载
        try:
            logger.info("尝试降级加载基础模型")
            self.load_model(model_path, None)  # 不加载LoRA
            return True
        except Exception as e2:
            logger.error(f"降级加载也失败: {e2}")
            return False

class InferenceMonitor:
    """推理监控器"""
    
    def __init__(self):
        self.metrics = {
            "total_requests": 0,
            "successful_requests": 0,
            "failed_requests": 0,
            "average_latency": 0.0,
            "model_usage": 0,
            "rules_usage": 0
        }
        self.latency_history = []
    
    def record_request(self, prompt: str, result: Dict, latency: float):
        """记录推理请求"""
        self.metrics["total_requests"] += 1
        self.latency_history.append(latency)
        
        if result["success"]:
            self.metrics["successful_requests"] += 1
            if result["method"] == "model":
                self.metrics["model_usage"] += 1
            else:
                self.metrics["rules_usage"] += 1
        else:
            self.metrics["failed_requests"] += 1
        
        # 更新平均延迟
        self.metrics["average_latency"] = sum(self.latency_history) / len(self.latency_history)
    
    def get_metrics(self) -> Dict:
        """获取监控指标"""
        return self.metrics.copy()
    
    def get_success_rate(self) -> float:
        """获取成功率"""
        if self.metrics["total_requests"] == 0:
            return 0.0
        return self.metrics["successful_requests"] / self.metrics["total_requests"]

def debug_inference(self, prompt: str) -> Dict[str, Any]:
    """调试推理过程"""
    debug_info = {
        "prompt": prompt,
        "model_loaded": self.model is not None,
        "tokenizer_loaded": self.tokenizer is not None,
        "generation_config": self.config["model"],
        "steps": []
    }
    
    # 步骤1：模型推理
    if self.model is not None:
        try:
            start_time = time.time()
            generated_text = self.generate_with_model(prompt)
            inference_time = time.time() - start_time
            
            debug_info["steps"].append({
                "step": "model_inference",
                "success": True,
                "time": inference_time,
                "output_length": len(generated_text),
                "raw_output": generated_text[:200] + "..." if len(generated_text) > 200 else generated_text
            })
            
            # 步骤2：JSON解析
            try:
                dsl = json.loads(generated_text)
                debug_info["steps"].append({
                    "step": "json_parsing",
                    "success": True,
                    "dsl_keys": list(dsl.keys()) if isinstance(dsl, dict) else []
                })
            except json.JSONDecodeError as e:
                debug_info["steps"].append({
                    "step": "json_parsing",
                    "success": False,
                    "error": str(e)
                })
                
        except Exception as e:
            debug_info["steps"].append({
                "step": "model_inference",
                "success": False,
                "error": str(e)
            })
    
    # 步骤3：规则回退
    try:
        start_time = time.time()
        rules_output = self.generate_with_rules(prompt)
        rules_time = time.time() - start_time
        
        debug_info["steps"].append({
            "step": "rules_fallback",
            "success": True,
            "time": rules_time,
            "output_length": len(rules_output)
        })
    except Exception as e:
        debug_info["steps"].append({
            "step": "rules_fallback",
            "success": False,
            "error": str(e)
        })
    
    return debug_info

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="UI生成推理")
    parser.add_argument("--prompt", type=str, required=True,
                       help="中文Prompt")
    parser.add_argument("--model_path", type=str, default="google/flan-t5-base",
                       help="模型路径")
    parser.add_argument("--lora_path", type=str, default=None,
                       help="LoRA权重路径")
    parser.add_argument("--output", type=str, required=True,
                       help="输出文件路径")
    parser.add_argument("--config", type=str, default="config/model_config.yaml",
                       help="配置文件路径")
    parser.add_argument("--use_rules", action="store_true",
                       help="强制使用规则生成")
    parser.add_argument("--debug", action="store_true",
                       help="调试模式")
    
    args = parser.parse_args()
    
    # 创建生成器
    generator = UIGenerator(args.config)
    
    # 加载模型（除非强制使用规则）
    if not args.use_rules:
        try:
            generator.load_model(args.model_path, args.lora_path)
        except Exception as e:
            logger.warning(f"模型加载失败: {e}，将使用规则生成")
    
    # 生成UI-DSL
    if args.debug:
        result = generator.debug_inference(args.prompt)
        print("调试信息:")
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        result = generator.generate(args.prompt, use_model=not args.use_rules)
    
    if result["success"]:
        # 保存结果
        output_path = Path(args.output)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(result["dsl"], f, ensure_ascii=False, indent=2)
        
        logger.info(f"UI-DSL已保存到: {output_path}")
        logger.info(f"生成方法: {result['method']}")
        
        # 打印生成的DSL
        print("\n生成的UI-DSL:")
        print(json.dumps(result["dsl"], ensure_ascii=False, indent=2))
    else:
        logger.error(f"生成失败: {result['error']}")
        exit(1)

if __name__ == "__main__":
    main()

# 推理性能优化建议
performance_recommendations = {
    "model_loading": {
        "use_quantization": True,
        "use_cache": True,
        "preload_models": True
    },
    "generation": {
        "batch_size": 4,
        "use_kv_cache": True,
        "optimize_sampling": True
    },
    "memory": {
        "use_fp16": True,
        "gradient_checkpointing": False,
        "device_map": "auto"
    }
}

# 质量保证策略
quality_assurance = {
    "validation": {
        "json_format_check": True,
        "content_quality_check": True,
        "fallback_mechanism": True
    },
    "monitoring": {
        "success_rate_tracking": True,
        "latency_monitoring": True,
        "error_logging": True
    },
    "testing": {
        "unit_tests": True,
        "integration_tests": True,
        "performance_tests": True
    }
}