import numpy as np
import json
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from sklearn.metrics import accuracy_score, f1_score

class TextUnderstandingEvaluator:
    def __init__(self, model_path, device="cuda"):
        """
        初始化评测器
        
        参数:
            model_path: 模型路径
            device: 运行设备
        """
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = AutoModelForCausalLM.from_pretrained(model_path).to(device)
        self.device = device
    
    def evaluate_multiple_choice(self, dataset_path, few_shot=0):
        """
        评测多项选择题能力
        
        参数:
            dataset_path: 数据集路径
            few_shot: few-shot示例数量
        """
        # 加载数据集
        dataset = load_dataset("json", data_files=dataset_path)["train"]
        
        correct_count = 0
        total_count = 0
        all_preds = []
        all_labels = []
        
        for item in dataset:
            # 构建提示
            prompt = self._construct_prompt(item, few_shot)
            
            # 获取模型预测
            prediction = self._get_model_prediction(prompt, item["choices"])
            
            # 记录结果
            if prediction == item["answer"]:
                correct_count += 1
            
            all_preds.append(prediction)
            all_labels.append(item["answer"])
            total_count += 1
        
        # 计算指标
        accuracy = correct_count / total_count
        f1 = f1_score(all_labels, all_preds, average="macro")
        
        return {
            "accuracy": accuracy,
            "f1": f1,
            "sample_count": total_count
        }
    
    def _construct_prompt(self, item, few_shot=0):
        """构建提示"""
        # 根据few-shot参数构建包含示例的提示
        prompt = f"问题：{item['question']}\n选项："
        
        for i, choice in enumerate(item["choices"]):
            prompt += f"\n{chr(65+i)}. {choice}"
        
        prompt += "\n请从选项中选择正确答案的选项字母。"
        return prompt
    
    def _get_model_prediction(self, prompt, choices):
        """获取模型预测"""
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
        
        with torch.no_grad():
            outputs = self.model.generate(
                inputs["input_ids"],
                max_new_tokens=5,
                num_return_sequences=1,
                temperature=0.1
            )
        
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        response = response[len(prompt):].strip()
        
        # 解析回答，查找选项字母
        for i, c in enumerate(['A', 'B', 'C', 'D', 'E']):
            if c in response and i < len(choices):
                return c
        
        return "A"  # 默认返回A

import nltk
from nltk.translate.bleu_score import corpus_bleu
from rouge_score import rouge_scorer
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

class TextGenerationEvaluator:
    def __init__(self, model_path, device="cuda"):
        """初始化文本生成评测器"""
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = AutoModelForCausalLM.from_pretrained(model_path).to(device)
        self.device = device
        self.rouge_scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2', 'rougeL'], use_stemmer=True)
    
    def evaluate_generation(self, prompts, references):
        """
        评估文本生成质量
        
        参数:
            prompts: 提示列表
            references: 参考回答列表(每个提示可能有多个参考回答)
        """
        generated_texts = []
        
        for prompt in prompts:
            generated_text = self._generate_text(prompt)
            generated_texts.append(generated_text)
        
        # 计算BLEU
        references_tokenized = [[ref.split()] for ref in references]
        generated_tokenized = [gen.split() for gen in generated_texts]
        bleu_score = corpus_bleu(references_tokenized, generated_tokenized)
        
        # 计算ROUGE
        rouge_scores = []
        for gen, ref in zip(generated_texts, references):
            score = self.rouge_scorer.score(gen, ref)
            rouge_scores.append(score)
        
        # 计算平均ROUGE分数
        avg_rouge1 = sum(score['rouge1'].fmeasure for score in rouge_scores) / len(rouge_scores)
        avg_rouge2 = sum(score['rouge2'].fmeasure for score in rouge_scores) / len(rouge_scores)
        avg_rougeL = sum(score['rougeL'].fmeasure for score in rouge_scores) / len(rouge_scores)
        
        # 计算困惑度
        perplexity = self._calculate_perplexity(generated_texts)
        
        return {
            "bleu": bleu_score,
            "rouge1": avg_rouge1,
            "rouge2": avg_rouge2,
            "rougeL": avg_rougeL,
            "perplexity": perplexity,
            "generated_samples": generated_texts[:5]  # 返回部分生成样本用于人工检查
        }
    
    def _generate_text(self, prompt, max_length=256):
        """生成文本"""
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
        
        with torch.no_grad():
            outputs = self.model.generate(
                inputs["input_ids"],
                max_new_tokens=max_length,
                num_return_sequences=1,
                temperature=0.7,
                top_p=0.9
            )
        
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response[len(prompt):].strip()
    
    def _calculate_perplexity(self, texts):
        """计算困惑度"""
        total_ppl = 0
        with torch.no_grad():
            for text in texts:
                inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
                outputs = self.model(inputs["input_ids"], labels=inputs["input_ids"])
                loss = outputs.loss
                ppl = torch.exp(loss)
                total_ppl += ppl.item()
        
        return total_ppl / len(texts)

def calculate_overall_score(scores, weights):
    """
    计算总体评分
    
    参数:
        scores: 各能力维度的得分字典
        weights: 各能力维度的权重字典
    
    返回:
        总体评分
    """
    overall_score = 0
    for category, score in scores.items():
        if category in weights:
            overall_score += score * weights[category]
    
    return overall_score

import matplotlib.pyplot as plt
import numpy as np

def plot_radar_chart(scores, categories):
    """
    绘制能力雷达图
    
    参数:
        scores: 各能力维度的得分列表
        categories: 各能力维度的名称列表
    """
    # 设置雷达图
    angles = np.linspace(0, 2*np.pi, len(categories), endpoint=False)
    scores = np.concatenate((scores, [scores[0]]))
    angles = np.concatenate((angles, [angles[0]]))
    categories = np.concatenate((categories, [categories[0]]))
    
    fig, ax = plt.subplots(figsize=(10, 10), subplot_kw=dict(polar=True))
    
    # 绘制得分线
    ax.plot(angles, scores, 'o-', linewidth=2)
    ax.fill(angles, scores, alpha=0.25)
    
    # 设置标签
    ax.set_thetagrids(angles[:-1] * 180/np.pi, categories[:-1])
    
    # 设置刻度
    ax.set_rlim(0, 100)
    ax.grid(True)
    
    plt.title('模型能力雷达图', size=20)
    plt.savefig('model_radar_chart.png', dpi=300, bbox_inches='tight')
    plt.show()