yutong.song

Full PLLM2.0 dump

4b1637c 11 months ago

24.4 kB

'''
多卡GPU inference:
# 使用所有可用GPU
python generating.py \
    --test_data_path /opt/tiger/PLLM1.0/dataset/processed_data/abstract_test.json \
    --output_path /opt/tiger/PLLM1.0/dataset/generated_answers.json \
    --base_model meta-llama/Llama-2-7b-hf \
    --embedding_model all-mpnet-base-v2 \
    --cluster_info_path /opt/tiger/PLLM1.0/clustering/cluster_loras/cluster_info.json \
    --lora_base_path /opt/tiger/PLLM1.0/clustering/cluster_loras \
    --max_length 256 \
    --temperature 0.3
# 指定使用4张GPU
python multi_gpu_generation.py \
    --num_gpus 4 \
    # ... 其他参数相同
'''

import json
import torch
import torch.multiprocessing as mp
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
from sklearn.metrics.pairwise import cosine_similarity
from tqdm import tqdm
import os
from typing import List, Dict, Any, Tuple
import argparse
import gc
import time
import threading
import queue
from concurrent.futures import ThreadPoolExecutor, as_completed
import fcntl
import copy

class MultiGPUAnswerGenerator:
    def __init__(self, 
                 base_model_name: str = "meta-llama/Llama-2-7b-hf",
                 embedding_model: str = "all-mpnet-base-v2",
                 cluster_info_path: str = "./cluster_loras/cluster_info.json",
                 lora_base_path: str = "./cluster_loras",
                 num_gpus: int = None):
        
        self.base_model_name = base_model_name
        self.embedding_model_name = embedding_model
        self.cluster_info_path = cluster_info_path
        self.lora_base_path = lora_base_path
        
        # 检测可用GPU数量
        if num_gpus is None:
            self.num_gpus = torch.cuda.device_count()
        else:
            self.num_gpus = min(num_gpus, torch.cuda.device_count())
        
        if self.num_gpus == 0:
            raise RuntimeError("没有检测到可用的GPU")
        
        print(f"🚀 使用 {self.num_gpus} 张GPU进行并行生成")
        for i in range(self.num_gpus):
            props = torch.cuda.get_device_properties(i)
            print(f"  GPU {i}: {props.name} ({props.total_memory/1024**3:.1f}GB)")
        
        # 加载embedding模型
        print(f"Loading embedding model: {embedding_model}")
        self.embedding_model = SentenceTransformer(embedding_model)
        
        # 加载聚类信息
        self.cluster_info = self.load_cluster_info()
        self.cluster_centers = self.compute_cluster_centers()
        
        # 实时写入相关
        self.write_lock = threading.Lock()
        self.results_queue = queue.Queue()
        self.writer_thread = None
        self.stop_writing = threading.Event()
    
    def load_cluster_info(self) -> Dict[str, Any]:
        """加载聚类信息"""
        print(f"Loading cluster info from {self.cluster_info_path}")
        with open(self.cluster_info_path, 'r', encoding='utf-8') as f:
            cluster_info = json.load(f)
        print(f"Loaded {cluster_info['n_clusters']} clusters")
        return cluster_info
    
    def compute_cluster_centers(self) -> np.ndarray:
        """计算聚类中心"""
        if 'cluster_centers' in self.cluster_info:
            centers = np.array(self.cluster_info['cluster_centers'])
            print(f"Using pre-computed cluster centers: {centers.shape}")
            return centers
        else:
            print("⚠️ No pre-computed cluster centers found")
            return None
    
    def extract_user_profile(self, user_data: Dict[str, Any]) -> str:
        """从用户数据中提取profile文本"""
        profile_text = ""
        if 'profile' in user_data and isinstance(user_data['profile'], list):
            profile_texts = []
            for profile_item in user_data['profile']:
                if isinstance(profile_item, dict):
                    if 'abstract' in profile_item:
                        profile_texts.append(profile_item['abstract'])
                    if 'title' in profile_item:
                        profile_texts.append(profile_item['title'])
            profile_text = ' '.join(profile_texts)
        return profile_text
    
    def find_nearest_cluster(self, user_embedding: np.ndarray) -> int:
        """找到最近的聚类"""
        if self.cluster_centers is None:
            return 0
        similarities = cosine_similarity([user_embedding], self.cluster_centers)[0]
        return np.argmax(similarities)
    
    def start_real_time_writer(self, output_path: str):
        """启动实时写入线程"""
        def writer_worker():
            written_results = []
            
            # 如果文件已存在，读取已有结果
            if os.path.exists(output_path):
                try:
                    with open(output_path, 'r', encoding='utf-8') as f:
                        written_results = json.load(f)
                    print(f"📂 恢复已有结果: {len(written_results)} 条记录")
                except:
                    written_results = []
            
            while not self.stop_writing.is_set() or not self.results_queue.empty():
                try:
                    # 等待新结果，超时1秒
                    new_results = self.results_queue.get(timeout=1.0)
                    written_results.extend(new_results)
                    
                    # 实时写入文件
                    with self.write_lock:
                        temp_path = output_path + '.tmp'
                        with open(temp_path, 'w', encoding='utf-8') as f:
                            json.dump(written_results, f, indent=2, ensure_ascii=False)
                        
                        # 原子性替换文件
                        os.replace(temp_path, output_path)
                    
                    print(f"💾 实时保存: 已写入 {len(written_results)} 条结果")
                    
                except queue.Empty:
                    continue
                except Exception as e:
                    print(f"❌ 写入错误: {e}")
        
        self.writer_thread = threading.Thread(target=writer_worker, daemon=True)
        self.writer_thread.start()
        print("✅ 实时写入线程已启动")
    
    def stop_real_time_writer(self):
        """停止实时写入线程"""
        if self.writer_thread:
            self.stop_writing.set()
            self.writer_thread.join(timeout=10)
            print("✅ 实时写入线程已停止")
    
    def load_model_on_gpu(self, cluster_id: int, gpu_id: int) -> Tuple[Any, Any]:
        """在指定GPU上加载模型"""
        cluster_dir = os.path.join(self.lora_base_path, f"cluster_{cluster_id}")
        
        if not os.path.exists(cluster_dir):
            print(f"⚠️ Cluster {cluster_id} directory not found")
            return None, None
        
        try:
            print(f"🔧 GPU {gpu_id}: 加载 cluster_{cluster_id} 模型...")
            
            # 设置当前GPU
            torch.cuda.set_device(gpu_id)
            
            # 加载tokenizer
            tokenizer = AutoTokenizer.from_pretrained(cluster_dir)
            if tokenizer.pad_token is None:
                tokenizer.pad_token = tokenizer.eos_token
            tokenizer.padding_side = "left"
            
            # 加载基础模型到指定GPU
            base_model = AutoModelForCausalLM.from_pretrained(
                self.base_model_name,
                torch_dtype=torch.float16,  # 使用float16节省内存
                device_map={f"": f"cuda:{gpu_id}"},  # 强制指定GPU
                low_cpu_mem_usage=True,
                max_memory={gpu_id: "70GB"}  # 为H100预留一些内存
            )
            
            # 加载LoRA适配器
            model = PeftModel.from_pretrained(base_model, cluster_dir)
            
            # 合并LoRA权重以提高推理速度
            try:
                model = model.merge_and_unload()
                print(f"✅ GPU {gpu_id}: LoRA权重已合并")
            except Exception as e:
                print(f"⚠️ GPU {gpu_id}: LoRA权重合并失败: {e}")
            
            model.eval()
            
            print(f"✅ GPU {gpu_id}: cluster_{cluster_id} 模型加载完成")
            return model, tokenizer
            
        except Exception as e:
            print(f"❌ GPU {gpu_id}: 加载模型失败: {e}")
            return None, None
    
    def generate_batch_on_gpu(self, samples: List[Dict], cluster_id: int, gpu_id: int,
                             max_length: int = 256, temperature: float = 0.3) -> List[Dict]:
        """在指定GPU上批量生成答案"""
        try:
            # 加载模型
            model, tokenizer = self.load_model_on_gpu(cluster_id, gpu_id)
            if model is None or tokenizer is None:
                return self._create_empty_results(samples, cluster_id)
            
            results = []
            batch_size = 4  # 每个GPU的batch size
            
            # 分批处理样本
            for i in range(0, len(samples), batch_size):
                batch = samples[i:i+batch_size]
                batch_queries = [s['query'] for s in batch]
                
                try:
                    # 编码输入
                    inputs = tokenizer(
                        batch_queries,
                        return_tensors="pt",
                        padding=True,
                        truncation=True,
                        max_length=128
                    )
                    
                    # 移动到指定GPU
                    device = f"cuda:{gpu_id}"
                    inputs = {k: v.to(device) for k, v in inputs.items()}
                    
                    # 生成答案
                    with torch.no_grad():
                        outputs = model.generate(
                            **inputs,
                            max_new_tokens=max_length,
                            temperature=temperature,
                            do_sample=True,
                            top_p=0.9,
                            top_k=50,
                            pad_token_id=tokenizer.eos_token_id,
                            eos_token_id=tokenizer.eos_token_id,
                            use_cache=True,
                            repetition_penalty=1.1
                        )
                    
                    # 解码答案
                    generated_texts = tokenizer.batch_decode(outputs, skip_special_tokens=True)
                    
                    # 构建结果
                    for j, sample in enumerate(batch):
                        answer = generated_texts[j][len(batch_queries[j]):].strip()
                        results.append({
                            'user_id': sample['user_id'],
                            'query': sample['query'],
                            'profile_text': sample['profile_text'],
                            'nearest_cluster': cluster_id,
                            'generated_answer': answer,
                            'ground_truth': sample['ground_truth'],
                            'gpu_id': gpu_id  # 记录使用的GPU
                        })
                    
                    # 每隔几个batch清理内存
                    if (i // batch_size) % 5 == 0:
                        torch.cuda.empty_cache()
                    
                except Exception as e:
                    print(f"❌ GPU {gpu_id} batch {i}-{i+batch_size} 生成失败: {e}")
                    # 添加空结果
                    for sample in batch:
                        results.append({
                            'user_id': sample['user_id'],
                            'query': sample['query'],
                            'profile_text': sample['profile_text'],
                            'nearest_cluster': cluster_id,
                            'generated_answer': "",
                            'ground_truth': sample['ground_truth'],
                            'gpu_id': gpu_id
                        })
            
            # 清理模型
            del model, tokenizer
            torch.cuda.empty_cache()
            
            print(f"✅ GPU {gpu_id}: 完成 {len(results)} 个样本的生成")
            return results
            
        except Exception as e:
            print(f"❌ GPU {gpu_id}: 处理失败: {e}")
            return self._create_empty_results(samples, cluster_id, gpu_id)
    
    def _create_empty_results(self, samples: List[Dict], cluster_id: int, gpu_id: int = -1) -> List[Dict]:
        """创建空结果"""
        return [{
            'user_id': s['user_id'],
            'query': s['query'],
            'profile_text': s['profile_text'],
            'nearest_cluster': cluster_id,
            'generated_answer': "",
            'ground_truth': s['ground_truth'],
            'gpu_id': gpu_id
        } for s in samples]
    
    def split_samples_for_gpus(self, samples: List[Dict], cluster_id: int) -> List[Tuple]:
        """将样本分配给多个GPU"""
        if not samples:
            return []
        
        # 计算每个GPU处理的样本数
        samples_per_gpu = len(samples) // self.num_gpus
        remainder = len(samples) % self.num_gpus
        
        gpu_tasks = []
        start_idx = 0
        
        for gpu_id in range(self.num_gpus):
            # 为前几个GPU分配额外的样本（如果有余数）
            current_batch_size = samples_per_gpu + (1 if gpu_id < remainder else 0)
            end_idx = start_idx + current_batch_size
            
            if start_idx < len(samples):
                gpu_samples = samples[start_idx:end_idx]
                gpu_tasks.append((gpu_samples, cluster_id, gpu_id))
                start_idx = end_idx
        
        return gpu_tasks
    
    def process_cluster_parallel(self, cluster_id: int, samples: List[Dict],
                                max_length: int = 256, temperature: float = 0.3) -> List[Dict]:
        """并行处理单个聚类的样本"""
        print(f"\n🚀 并行处理 cluster_{cluster_id}: {len(samples)} 个样本")
        
        # 将样本分配给多个GPU
        gpu_tasks = self.split_samples_for_gpus(samples, cluster_id)
        
        if not gpu_tasks:
            return []
        
        print(f"📋 任务分配:")
        for i, (gpu_samples, _, gpu_id) in enumerate(gpu_tasks):
            print(f"  GPU {gpu_id}: {len(gpu_samples)} 个样本")
        
        all_results = []
        
        # 使用线程池并行处理
        with ThreadPoolExecutor(max_workers=self.num_gpus) as executor:
            future_to_gpu = {}
            
            # 提交任务到线程池
            for gpu_samples, cluster_id, gpu_id in gpu_tasks:
                future = executor.submit(
                    self.generate_batch_on_gpu,
                    gpu_samples, cluster_id, gpu_id, max_length, temperature
                )
                future_to_gpu[future] = gpu_id
            
            # 收集结果
            for future in as_completed(future_to_gpu):
                gpu_id = future_to_gpu[future]
                try:
                    results = future.result()
                    all_results.extend(results)
                    
                    # 实时将结果放入队列
                    if results:
                        self.results_queue.put(results)
                    
                    print(f"✅ GPU {gpu_id} 完成处理")
                    
                except Exception as e:
                    print(f"❌ GPU {gpu_id} 处理异常: {e}")
        
        print(f"✅ Cluster {cluster_id} 并行处理完成: {len(all_results)} 个结果")
        return all_results
    
    def convert_numpy_types(self, obj):
        """转换numpy类型为Python原生类型"""
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, dict):
            return {key: self.convert_numpy_types(value) for key, value in obj.items()}
        elif isinstance(obj, list):
            return [self.convert_numpy_types(item) for item in obj]
        else:
            return obj
    
    def process_test_data(self, test_data_path: str, output_path: str,
                         max_length: int = 256, temperature: float = 0.3):
        """处理测试数据的主函数"""
        print(f"📂 处理测试数据: {test_data_path}")
        
        # 启动实时写入线程
        self.start_real_time_writer(output_path)
        
        try:
            # 加载测试数据
            with open(test_data_path, 'r', encoding='utf-8') as f:
                test_data = json.load(f)
            
            # 加载或计算聚类分配
            cluster_to_samples = self._load_cluster_assignments(test_data)
            
            # 按聚类处理
            total_processed = 0
            start_time = time.time()
            
            for cluster_id, samples in cluster_to_samples.items():
                print(f"\n{'='*60}")
                print(f"处理聚类 {cluster_id}: {len(samples)} 个样本")
                
                cluster_start = time.time()
                
                # 并行处理聚类样本
                cluster_results = self.process_cluster_parallel(
                    cluster_id, samples, max_length, temperature
                )
                
                total_processed += len(cluster_results)
                cluster_time = time.time() - cluster_start
                
                print(f"✅ Cluster {cluster_id} 完成，耗时: {cluster_time:.1f}s")
                print(f"📊 处理速度: {len(samples)/cluster_time:.1f} 样本/秒")
                print(f"📈 总进度: {total_processed}/{len(test_data)} ({total_processed/len(test_data)*100:.1f}%)")
                
                # 清理GPU内存
                for gpu_id in range(self.num_gpus):
                    torch.cuda.set_device(gpu_id)
                    torch.cuda.empty_cache()
                gc.collect()
            
            # 等待写入完成
            print("\n⏳ 等待所有结果写入完成...")
            time.sleep(2)  # 给写入线程一些时间
            
            total_time = time.time() - start_time
            print(f"\n🎉 全部处理完成！")
            print(f"⏱️ 总耗时: {total_time:.1f}s")
            print(f"⚡ 平均速度: {total_processed/total_time:.1f} 样本/秒")
            print(f"💾 结果保存到: {output_path}")
            
        except Exception as e:
            print(f"❌ 处理过程中出错: {e}")
            raise
        finally:
            # 停止实时写入线程
            self.stop_real_time_writer()
    
    def _load_cluster_assignments(self, test_data):
        """加载或计算聚类分配"""
        assignment_path = "cluster_assignment.json"
        cluster_to_samples = {}
        
        if os.path.exists(assignment_path):
            print(f"📂 加载已有的聚类分配: {assignment_path}")
            with open(assignment_path, "r", encoding="utf-8") as f:
                assignments = {item["user_id"]: item["nearest_cluster"] for item in json.load(f)}
            
            for i, item in enumerate(test_data):
                user_id = item.get('name', f'user_{i}')
                query = item.get('input', '')
                if not query:
                    continue
                    
                profile_text = self.extract_user_profile(item)
                nearest_cluster = assignments.get(user_id, 0)
                
                if nearest_cluster not in cluster_to_samples:
                    cluster_to_samples[nearest_cluster] = []
                cluster_to_samples[nearest_cluster].append({
                    'user_id': user_id,
                    'query': query,
                    'profile_text': profile_text,
                    'ground_truth': item.get('output', '')
                })
        else:
            print("🔄 计算聚类分配...")
            assignments_list = []
            
            for i, item in enumerate(tqdm(test_data, desc="Computing assignments")):
                user_id = item.get('name', f'user_{i}')
                query = item.get('input', '')
                if not query:
                    continue
                    
                profile_text = self.extract_user_profile(item)
                user_embedding = self.embedding_model.encode([profile_text])[0]
                nearest_cluster = self.find_nearest_cluster(user_embedding)
                assignments_list.append({"user_id": user_id, "nearest_cluster": int(nearest_cluster)})
                
                if nearest_cluster not in cluster_to_samples:
                    cluster_to_samples[nearest_cluster] = []
                cluster_to_samples[nearest_cluster].append({
                    'user_id': user_id,
                    'query': query,
                    'profile_text': profile_text,
                    'ground_truth': item.get('output', '')
                })
            
            # 保存聚类分配结果
            with open(assignment_path, "w", encoding="utf-8") as f:
                json.dump(assignments_list, f, ensure_ascii=False, indent=2)
        
        print(f"📊 聚类分布:")
        for cluster_id, samples in cluster_to_samples.items():
            print(f"  Cluster {cluster_id}: {len(samples)} 个样本")
        
        return cluster_to_samples

def main():
    # 设置多进程启动方法
    mp.set_start_method('spawn', force=True)
    
    parser = argparse.ArgumentParser(description='多GPU并行聚类LoRA生成器')
    parser.add_argument('--test_data_path', type=str, required=True,
                       help='测试数据文件路径')
    parser.add_argument('--output_path', type=str, default='./generated_answers.json',
                       help='输出文件路径')
    parser.add_argument('--base_model', type=str, default='meta-llama/Llama-2-7b-hf',
                       help='基础模型名称')
    parser.add_argument('--embedding_model', type=str, default='all-mpnet-base-v2',
                       help='嵌入模型名称')
    parser.add_argument('--cluster_info_path', type=str, default='./cluster_loras/cluster_info.json',
                       help='聚类信息文件路径')
    parser.add_argument('--lora_base_path', type=str, default='./cluster_loras',
                       help='LoRA模型基础路径')
    parser.add_argument('--max_length', type=int, default=256,
                       help='生成答案的最大长度')
    parser.add_argument('--temperature', type=float, default=0.3,
                       help='生成温度')
    parser.add_argument('--num_gpus', type=int, default=None,
                       help='使用的GPU数量（默认使用所有可用GPU）')
    
    args = parser.parse_args()
    
    # 创建输出目录
    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
    
    try:
        # 初始化多GPU生成器
        generator = MultiGPUAnswerGenerator(
            base_model_name=args.base_model,
            embedding_model=args.embedding_model,
            cluster_info_path=args.cluster_info_path,
            lora_base_path=args.lora_base_path,
            num_gpus=args.num_gpus
        )
        
        # 处理测试数据
        generator.process_test_data(
            args.test_data_path,
            args.output_path,
            max_length=args.max_length,
            temperature=args.temperature
        )
        
        print(f"\n✅ 多GPU并行生成完成!")
        print(f"📁 结果文件: {args.output_path}")
        
    except Exception as e:
        print(f"❌ 程序执行失败: {e}")
        # 清理所有GPU内存
        for gpu_id in range(torch.cuda.device_count()):
            torch.cuda.set_device(gpu_id)
            torch.cuda.empty_cache()
        raise

if __name__ == "__main__":
    main()


	'''
	多卡GPU inference:
	# 使用所有可用GPU
	python generating.py \
	--test_data_path /opt/tiger/PLLM1.0/dataset/processed_data/abstract_test.json \
	--output_path /opt/tiger/PLLM1.0/dataset/generated_answers.json \
	--base_model meta-llama/Llama-2-7b-hf \
	--embedding_model all-mpnet-base-v2 \
	--cluster_info_path /opt/tiger/PLLM1.0/clustering/cluster_loras/cluster_info.json \
	--lora_base_path /opt/tiger/PLLM1.0/clustering/cluster_loras \
	--max_length 256 \
	--temperature 0.3

	# 指定使用4张GPU
	python multi_gpu_generation.py \
	--num_gpus 4 \
	# ... 其他参数相同

	'''

	import json
	import torch
	import torch.multiprocessing as mp
	import numpy as np
	from sentence_transformers import SentenceTransformer
	from transformers import AutoTokenizer, AutoModelForCausalLM
	from peft import PeftModel
	from sklearn.metrics.pairwise import cosine_similarity
	from tqdm import tqdm
	import os
	from typing import List, Dict, Any, Tuple
	import argparse
	import gc
	import time
	import threading
	import queue
	from concurrent.futures import ThreadPoolExecutor, as_completed
	import fcntl
	import copy

	class MultiGPUAnswerGenerator:
	def __init__(self,
	base_model_name: str = "meta-llama/Llama-2-7b-hf",
	embedding_model: str = "all-mpnet-base-v2",
	cluster_info_path: str = "./cluster_loras/cluster_info.json",
	lora_base_path: str = "./cluster_loras",
	num_gpus: int = None):

	self.base_model_name = base_model_name
	self.embedding_model_name = embedding_model
	self.cluster_info_path = cluster_info_path
	self.lora_base_path = lora_base_path

	# 检测可用GPU数量
	if num_gpus is None:
	self.num_gpus = torch.cuda.device_count()
	else:
	self.num_gpus = min(num_gpus, torch.cuda.device_count())

	if self.num_gpus == 0:
	raise RuntimeError("没有检测到可用的GPU")

	print(f"🚀 使用 {self.num_gpus} 张GPU进行并行生成")
	for i in range(self.num_gpus):
	props = torch.cuda.get_device_properties(i)
	print(f" GPU {i}: {props.name} ({props.total_memory/1024**3:.1f}GB)")

	# 加载embedding模型
	print(f"Loading embedding model: {embedding_model}")
	self.embedding_model = SentenceTransformer(embedding_model)

	# 加载聚类信息
	self.cluster_info = self.load_cluster_info()
	self.cluster_centers = self.compute_cluster_centers()

	# 实时写入相关
	self.write_lock = threading.Lock()
	self.results_queue = queue.Queue()
	self.writer_thread = None
	self.stop_writing = threading.Event()

	def load_cluster_info(self) -> Dict[str, Any]:
	"""加载聚类信息"""
	print(f"Loading cluster info from {self.cluster_info_path}")
	with open(self.cluster_info_path, 'r', encoding='utf-8') as f:
	cluster_info = json.load(f)
	print(f"Loaded {cluster_info['n_clusters']} clusters")
	return cluster_info

	def compute_cluster_centers(self) -> np.ndarray:
	"""计算聚类中心"""
	if 'cluster_centers' in self.cluster_info:
	centers = np.array(self.cluster_info['cluster_centers'])
	print(f"Using pre-computed cluster centers: {centers.shape}")
	return centers
	else:
	print("⚠️ No pre-computed cluster centers found")
	return None

	def extract_user_profile(self, user_data: Dict[str, Any]) -> str:
	"""从用户数据中提取profile文本"""
	profile_text = ""
	if 'profile' in user_data and isinstance(user_data['profile'], list):
	profile_texts = []
	for profile_item in user_data['profile']:
	if isinstance(profile_item, dict):
	if 'abstract' in profile_item:
	profile_texts.append(profile_item['abstract'])
	if 'title' in profile_item:
	profile_texts.append(profile_item['title'])
	profile_text = ' '.join(profile_texts)
	return profile_text

	def find_nearest_cluster(self, user_embedding: np.ndarray) -> int:
	"""找到最近的聚类"""
	if self.cluster_centers is None:
	return 0
	similarities = cosine_similarity([user_embedding], self.cluster_centers)[0]
	return np.argmax(similarities)

	def start_real_time_writer(self, output_path: str):
	"""启动实时写入线程"""
	def writer_worker():
	written_results = []

	# 如果文件已存在，读取已有结果
	if os.path.exists(output_path):
	try:
	with open(output_path, 'r', encoding='utf-8') as f:
	written_results = json.load(f)
	print(f"📂 恢复已有结果: {len(written_results)} 条记录")
	except:
	written_results = []

	while not self.stop_writing.is_set() or not self.results_queue.empty():
	try:
	# 等待新结果，超时1秒
	new_results = self.results_queue.get(timeout=1.0)
	written_results.extend(new_results)

	# 实时写入文件
	with self.write_lock:
	temp_path = output_path + '.tmp'
	with open(temp_path, 'w', encoding='utf-8') as f:
	json.dump(written_results, f, indent=2, ensure_ascii=False)

	# 原子性替换文件
	os.replace(temp_path, output_path)

	print(f"💾 实时保存: 已写入 {len(written_results)} 条结果")

	except queue.Empty:
	continue
	except Exception as e:
	print(f"❌ 写入错误: {e}")

	self.writer_thread = threading.Thread(target=writer_worker, daemon=True)
	self.writer_thread.start()
	print("✅ 实时写入线程已启动")

	def stop_real_time_writer(self):
	"""停止实时写入线程"""
	if self.writer_thread:
	self.stop_writing.set()
	self.writer_thread.join(timeout=10)
	print("✅ 实时写入线程已停止")

	def load_model_on_gpu(self, cluster_id: int, gpu_id: int) -> Tuple[Any, Any]:
	"""在指定GPU上加载模型"""
	cluster_dir = os.path.join(self.lora_base_path, f"cluster_{cluster_id}")

	if not os.path.exists(cluster_dir):
	print(f"⚠️ Cluster {cluster_id} directory not found")
	return None, None

	try:
	print(f"🔧 GPU {gpu_id}: 加载 cluster_{cluster_id} 模型...")

	# 设置当前GPU
	torch.cuda.set_device(gpu_id)

	# 加载tokenizer
	tokenizer = AutoTokenizer.from_pretrained(cluster_dir)
	if tokenizer.pad_token is None:
	tokenizer.pad_token = tokenizer.eos_token
	tokenizer.padding_side = "left"

	# 加载基础模型到指定GPU
	base_model = AutoModelForCausalLM.from_pretrained(
	self.base_model_name,
	torch_dtype=torch.float16, # 使用float16节省内存
	device_map={f"": f"cuda:{gpu_id}"}, # 强制指定GPU
	low_cpu_mem_usage=True,
	max_memory={gpu_id: "70GB"} # 为H100预留一些内存
	)

	# 加载LoRA适配器
	model = PeftModel.from_pretrained(base_model, cluster_dir)

	# 合并LoRA权重以提高推理速度
	try:
	model = model.merge_and_unload()
	print(f"✅ GPU {gpu_id}: LoRA权重已合并")
	except Exception as e:
	print(f"⚠️ GPU {gpu_id}: LoRA权重合并失败: {e}")

	model.eval()

	print(f"✅ GPU {gpu_id}: cluster_{cluster_id} 模型加载完成")
	return model, tokenizer

	except Exception as e:
	print(f"❌ GPU {gpu_id}: 加载模型失败: {e}")
	return None, None

	def generate_batch_on_gpu(self, samples: List[Dict], cluster_id: int, gpu_id: int,
	max_length: int = 256, temperature: float = 0.3) -> List[Dict]:
	"""在指定GPU上批量生成答案"""
	try:
	# 加载模型
	model, tokenizer = self.load_model_on_gpu(cluster_id, gpu_id)
	if model is None or tokenizer is None:
	return self._create_empty_results(samples, cluster_id)

	results = []
	batch_size = 4 # 每个GPU的batch size

	# 分批处理样本
	for i in range(0, len(samples), batch_size):
	batch = samples[i:i+batch_size]
	batch_queries = [s['query'] for s in batch]

	try:
	# 编码输入
	inputs = tokenizer(
	batch_queries,
	return_tensors="pt",
	padding=True,
	truncation=True,
	max_length=128
	)

	# 移动到指定GPU
	device = f"cuda:{gpu_id}"
	inputs = {k: v.to(device) for k, v in inputs.items()}

	# 生成答案
	with torch.no_grad():
	outputs = model.generate(
	**inputs,
	max_new_tokens=max_length,
	temperature=temperature,
	do_sample=True,
	top_p=0.9,
	top_k=50,
	pad_token_id=tokenizer.eos_token_id,
	eos_token_id=tokenizer.eos_token_id,
	use_cache=True,
	repetition_penalty=1.1
	)

	# 解码答案
	generated_texts = tokenizer.batch_decode(outputs, skip_special_tokens=True)

	# 构建结果
	for j, sample in enumerate(batch):
	answer = generated_texts[j][len(batch_queries[j]):].strip()
	results.append({
	'user_id': sample['user_id'],
	'query': sample['query'],
	'profile_text': sample['profile_text'],
	'nearest_cluster': cluster_id,
	'generated_answer': answer,
	'ground_truth': sample['ground_truth'],
	'gpu_id': gpu_id # 记录使用的GPU
	})

	# 每隔几个batch清理内存
	if (i // batch_size) % 5 == 0:
	torch.cuda.empty_cache()

	except Exception as e:
	print(f"❌ GPU {gpu_id} batch {i}-{i+batch_size} 生成失败: {e}")
	# 添加空结果
	for sample in batch:
	results.append({
	'user_id': sample['user_id'],
	'query': sample['query'],
	'profile_text': sample['profile_text'],
	'nearest_cluster': cluster_id,
	'generated_answer': "",
	'ground_truth': sample['ground_truth'],
	'gpu_id': gpu_id
	})

	# 清理模型
	del model, tokenizer
	torch.cuda.empty_cache()

	print(f"✅ GPU {gpu_id}: 完成 {len(results)} 个样本的生成")
	return results

	except Exception as e:
	print(f"❌ GPU {gpu_id}: 处理失败: {e}")
	return self._create_empty_results(samples, cluster_id, gpu_id)

	def _create_empty_results(self, samples: List[Dict], cluster_id: int, gpu_id: int = -1) -> List[Dict]:
	"""创建空结果"""
	return [{
	'user_id': s['user_id'],
	'query': s['query'],
	'profile_text': s['profile_text'],
	'nearest_cluster': cluster_id,
	'generated_answer': "",
	'ground_truth': s['ground_truth'],
	'gpu_id': gpu_id
	} for s in samples]

	def split_samples_for_gpus(self, samples: List[Dict], cluster_id: int) -> List[Tuple]:
	"""将样本分配给多个GPU"""
	if not samples:
	return []

	# 计算每个GPU处理的样本数
	samples_per_gpu = len(samples) // self.num_gpus
	remainder = len(samples) % self.num_gpus

	gpu_tasks = []
	start_idx = 0

	for gpu_id in range(self.num_gpus):
	# 为前几个GPU分配额外的样本（如果有余数）
	current_batch_size = samples_per_gpu + (1 if gpu_id < remainder else 0)
	end_idx = start_idx + current_batch_size

	if start_idx < len(samples):
	gpu_samples = samples[start_idx:end_idx]
	gpu_tasks.append((gpu_samples, cluster_id, gpu_id))
	start_idx = end_idx

	return gpu_tasks

	def process_cluster_parallel(self, cluster_id: int, samples: List[Dict],
	max_length: int = 256, temperature: float = 0.3) -> List[Dict]:
	"""并行处理单个聚类的样本"""
	print(f"\n🚀 并行处理 cluster_{cluster_id}: {len(samples)} 个样本")

	# 将样本分配给多个GPU
	gpu_tasks = self.split_samples_for_gpus(samples, cluster_id)

	if not gpu_tasks:
	return []

	print(f"📋 任务分配:")
	for i, (gpu_samples, _, gpu_id) in enumerate(gpu_tasks):
	print(f" GPU {gpu_id}: {len(gpu_samples)} 个样本")

	all_results = []

	# 使用线程池并行处理
	with ThreadPoolExecutor(max_workers=self.num_gpus) as executor:
	future_to_gpu = {}

	# 提交任务到线程池
	for gpu_samples, cluster_id, gpu_id in gpu_tasks:
	future = executor.submit(
	self.generate_batch_on_gpu,
	gpu_samples, cluster_id, gpu_id, max_length, temperature
	)
	future_to_gpu[future] = gpu_id

	# 收集结果
	for future in as_completed(future_to_gpu):
	gpu_id = future_to_gpu[future]
	try:
	results = future.result()
	all_results.extend(results)

	# 实时将结果放入队列
	if results:
	self.results_queue.put(results)

	print(f"✅ GPU {gpu_id} 完成处理")

	except Exception as e:
	print(f"❌ GPU {gpu_id} 处理异常: {e}")

	print(f"✅ Cluster {cluster_id} 并行处理完成: {len(all_results)} 个结果")
	return all_results

	def convert_numpy_types(self, obj):
	"""转换numpy类型为Python原生类型"""
	if isinstance(obj, np.integer):
	return int(obj)
	elif isinstance(obj, np.floating):
	return float(obj)
	elif isinstance(obj, np.ndarray):
	return obj.tolist()
	elif isinstance(obj, dict):
	return {key: self.convert_numpy_types(value) for key, value in obj.items()}
	elif isinstance(obj, list):
	return [self.convert_numpy_types(item) for item in obj]
	else:
	return obj

	def process_test_data(self, test_data_path: str, output_path: str,
	max_length: int = 256, temperature: float = 0.3):
	"""处理测试数据的主函数"""
	print(f"📂 处理测试数据: {test_data_path}")

	# 启动实时写入线程
	self.start_real_time_writer(output_path)

	try:
	# 加载测试数据
	with open(test_data_path, 'r', encoding='utf-8') as f:
	test_data = json.load(f)

	# 加载或计算聚类分配
	cluster_to_samples = self._load_cluster_assignments(test_data)

	# 按聚类处理
	total_processed = 0
	start_time = time.time()

	for cluster_id, samples in cluster_to_samples.items():
	print(f"\n{'='*60}")
	print(f"处理聚类 {cluster_id}: {len(samples)} 个样本")

	cluster_start = time.time()

	# 并行处理聚类样本
	cluster_results = self.process_cluster_parallel(
	cluster_id, samples, max_length, temperature
	)

	total_processed += len(cluster_results)
	cluster_time = time.time() - cluster_start

	print(f"✅ Cluster {cluster_id} 完成，耗时: {cluster_time:.1f}s")
	print(f"📊 处理速度: {len(samples)/cluster_time:.1f} 样本/秒")
	print(f"📈 总进度: {total_processed}/{len(test_data)} ({total_processed/len(test_data)*100:.1f}%)")

	# 清理GPU内存
	for gpu_id in range(self.num_gpus):
	torch.cuda.set_device(gpu_id)
	torch.cuda.empty_cache()
	gc.collect()

	# 等待写入完成
	print("\n⏳ 等待所有结果写入完成...")
	time.sleep(2) # 给写入线程一些时间

	total_time = time.time() - start_time
	print(f"\n🎉 全部处理完成！")
	print(f"⏱️ 总耗时: {total_time:.1f}s")
	print(f"⚡ 平均速度: {total_processed/total_time:.1f} 样本/秒")
	print(f"💾 结果保存到: {output_path}")

	except Exception as e:
	print(f"❌ 处理过程中出错: {e}")
	raise
	finally:
	# 停止实时写入线程
	self.stop_real_time_writer()

	def _load_cluster_assignments(self, test_data):
	"""加载或计算聚类分配"""
	assignment_path = "cluster_assignment.json"
	cluster_to_samples = {}

	if os.path.exists(assignment_path):
	print(f"📂 加载已有的聚类分配: {assignment_path}")
	with open(assignment_path, "r", encoding="utf-8") as f:
	assignments = {item["user_id"]: item["nearest_cluster"] for item in json.load(f)}

	for i, item in enumerate(test_data):
	user_id = item.get('name', f'user_{i}')
	query = item.get('input', '')
	if not query:
	continue

	profile_text = self.extract_user_profile(item)
	nearest_cluster = assignments.get(user_id, 0)

	if nearest_cluster not in cluster_to_samples:
	cluster_to_samples[nearest_cluster] = []
	cluster_to_samples[nearest_cluster].append({
	'user_id': user_id,
	'query': query,
	'profile_text': profile_text,
	'ground_truth': item.get('output', '')
	})
	else:
	print("🔄 计算聚类分配...")
	assignments_list = []

	for i, item in enumerate(tqdm(test_data, desc="Computing assignments")):
	user_id = item.get('name', f'user_{i}')
	query = item.get('input', '')
	if not query:
	continue

	profile_text = self.extract_user_profile(item)
	user_embedding = self.embedding_model.encode([profile_text])[0]
	nearest_cluster = self.find_nearest_cluster(user_embedding)
	assignments_list.append({"user_id": user_id, "nearest_cluster": int(nearest_cluster)})

	if nearest_cluster not in cluster_to_samples:
	cluster_to_samples[nearest_cluster] = []
	cluster_to_samples[nearest_cluster].append({
	'user_id': user_id,
	'query': query,
	'profile_text': profile_text,
	'ground_truth': item.get('output', '')
	})

	# 保存聚类分配结果
	with open(assignment_path, "w", encoding="utf-8") as f:
	json.dump(assignments_list, f, ensure_ascii=False, indent=2)

	print(f"📊 聚类分布:")
	for cluster_id, samples in cluster_to_samples.items():
	print(f" Cluster {cluster_id}: {len(samples)} 个样本")

	return cluster_to_samples

	def main():
	# 设置多进程启动方法
	mp.set_start_method('spawn', force=True)

	parser = argparse.ArgumentParser(description='多GPU并行聚类LoRA生成器')
	parser.add_argument('--test_data_path', type=str, required=True,
	help='测试数据文件路径')
	parser.add_argument('--output_path', type=str, default='./generated_answers.json',
	help='输出文件路径')
	parser.add_argument('--base_model', type=str, default='meta-llama/Llama-2-7b-hf',
	help='基础模型名称')
	parser.add_argument('--embedding_model', type=str, default='all-mpnet-base-v2',
	help='嵌入模型名称')
	parser.add_argument('--cluster_info_path', type=str, default='./cluster_loras/cluster_info.json',
	help='聚类信息文件路径')
	parser.add_argument('--lora_base_path', type=str, default='./cluster_loras',
	help='LoRA模型基础路径')
	parser.add_argument('--max_length', type=int, default=256,
	help='生成答案的最大长度')
	parser.add_argument('--temperature', type=float, default=0.3,
	help='生成温度')
	parser.add_argument('--num_gpus', type=int, default=None,
	help='使用的GPU数量（默认使用所有可用GPU）')

	args = parser.parse_args()

	# 创建输出目录
	os.makedirs(os.path.dirname(args.output_path), exist_ok=True)

	try:
	# 初始化多GPU生成器
	generator = MultiGPUAnswerGenerator(
	base_model_name=args.base_model,
	embedding_model=args.embedding_model,
	cluster_info_path=args.cluster_info_path,
	lora_base_path=args.lora_base_path,
	num_gpus=args.num_gpus
	)

	# 处理测试数据
	generator.process_test_data(
	args.test_data_path,
	args.output_path,
	max_length=args.max_length,
	temperature=args.temperature
	)

	print(f"\n✅ 多GPU并行生成完成!")
	print(f"📁 结果文件: {args.output_path}")

	except Exception as e:
	print(f"❌ 程序执行失败: {e}")
	# 清理所有GPU内存
	for gpu_id in range(torch.cuda.device_count()):
	torch.cuda.set_device(gpu_id)
	torch.cuda.empty_cache()
	raise

	if __name__ == "__main__":
	main()