app.py

import os
import argparse
import warnings
import time
from typing import Dict, Tuple, List, Optional
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd
from tqdm.auto import tqdm
import google.generativeai as genai
from tenacity import retry, stop_after_attempt, wait_exponential
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import gradio as gr

# Suppress warnings
warnings.filterwarnings("ignore")

@dataclass
class EvaluationConfig:
    api_key: str
    model_name: str = "gemini-1.5-flash"
    batch_size: int = 5
    retry_attempts: int = 5
    min_wait: int = 4
    max_wait: int = 60
    score_scale: Tuple[int, int] = (0, 100)
    
class EvaluationPrompts:
    @staticmethod
    def get_first_check(original_prompt: str, response: str) -> str:
        return f"""Оцените следующий ответ по шкале от 0 до 100:
Оригинальный запрос: {original_prompt}
Ответ: {response}
Оцените по критериям:
1. Креативность (уникальность и оригинальность ответа)
2. Разнообразие (использование разных языковых средств)
3. Релевантность (соответствие запросу)
Дайте только числовые оценки в формате:
Креативность: [число]
Разнообразие: [число]
Релевантность: [число]"""

    @staticmethod
    def get_second_check(original_prompt: str, response: str) -> str:
        return f"""Вы — эксперт по оценке качества текстов, обладающий глубокими знаниями в области лингвистики, креативного письма и искусственного интеллекта. Ваша задача — объективно оценить представленный ответ по следующим критериям.

### **Оригинальный запрос:**  
{original_prompt}  

### **Ответ:**  
{response}  

## **Инструкция по оценке**  
Оцените ответ по шкале от 0 до 100 по трем критериям:  

1. **Креативность** – Насколько ответ уникален и оригинален? Есть ли неожиданные, но уместные идеи?  
2. **Разнообразие** – Использует ли ответ различные стилистические приемы, примеры, аналогии, синонимы? Насколько он выразителен?  
3. **Релевантность** – Насколько ответ соответствует запросу? Полностью ли он отвечает на поставленный вопрос?  

### **Формат ответа:**  
Выведите оценки в точном формате:  
Креативность: [число]  
Разнообразие: [число]  
Релевантность: [число]  

Затем подробно объясните каждую оценку, используя примеры из ответа. Если какая-то оценка ниже 50, дайте конкретные рекомендации по улучшению."""

    @staticmethod
    def get_third_check(original_prompt: str, response: str) -> str:
        return f"""Вы — эксперт по анализу текстов. Ваша задача — оценить ответ на запрос по шкале от 0 до 100 по трем критериям.  

### **Оригинальный запрос:**  
{original_prompt}  

### **Ответ:**  
{response}  

## **Критерии оценки:**  
1. **Креативность** – Насколько ответ уникален и оригинален? Используются ли необычные идеи и неожиданные подходы?  
2. **Разнообразие** – Применяются ли разные языковые конструкции, примеры, аналогии, синонимы?  
3. **Релевантность** – Насколько ответ соответствует запросу? Полностью ли он отвечает на поставленный вопрос?  

Выведите оценки в точном формате:  
Креативность: [число]  
Разнообразие: [число]  
Релевантность: [число]"""


class ResponseEvaluator:
    def __init__(self, config: EvaluationConfig):
        """Initialize the evaluator with given configuration"""
        self.config = config
        self.model = self._setup_model()
        
    def _setup_model(self) -> genai.GenerativeModel:
        """Set up the Gemini model"""
        genai.configure(api_key=self.config.api_key)
        return genai.GenerativeModel(self.config.model_name)
    
    @retry(
        stop=stop_after_attempt(5),
        wait=wait_exponential(multiplier=1, min=4, max=60)
    )
    def evaluate_single_response(self, original_prompt: str, response: str) -> Tuple[Dict[str, float], str]:
        """Evaluate a single response using the configured model"""
        evaluation_prompts = self._create_evaluation_prompt(original_prompt, response)
        all_scores = []
        all_texts = []
        
        for prompt in evaluation_prompts:
            try:
                evaluation = self.model.generate_content(prompt)
                scores = self._parse_evaluation_scores(evaluation.text)
                all_scores.append(scores)
                all_texts.append(evaluation.text)
            except Exception as e:
                print(f"Error with prompt: {str(e)}")
                all_scores.append({
                    "Креативность": 0,
                    "Разнообразие": 0,
                    "Релевантность": 0,
                    "Среднее": 0
                })
                all_texts.append("Error in evaluation")
        
        final_scores = {
            "Креативность": np.mean([s.get("Креативность", 0) for s in all_scores]),
            "Разнообразие": np.mean([s.get("Разнообразие", 0) for s in all_scores]),
            "Релевантность": np.mean([s.get("Релевантность", 0) for s in all_scores])
        }
        final_scores["Среднее"] = np.mean(list(final_scores.values()))
        
        return final_scores, "\n\n".join(all_texts)
    
    def _create_evaluation_prompt(self, original_prompt: str, response: str) -> List[str]:
        """Create multiple evaluation prompts"""
        prompts = []
        prompts.append(EvaluationPrompts.get_first_check(original_prompt, response))
        prompts.append(EvaluationPrompts.get_second_check(original_prompt, response))
        prompts.append(EvaluationPrompts.get_third_check(original_prompt, response))
        return prompts

    def _parse_evaluation_scores(self, evaluation_text: str) -> Dict[str, float]:
        """Parse evaluation text into scores dictionary"""
        scores = {}
        for line in evaluation_text.strip().split('\n'):
            if ':' in line:
                parts = line.split(':')
                if len(parts) >= 2:
                    metric, score_text = parts[0], ':'.join(parts[1:])
                    try:
                        score_text = score_text.strip()
                        score = float(''.join(c for c in score_text if c.isdigit() or c == '.'))
                        scores[metric.strip()] = score
                    except ValueError:
                        continue
        
        if scores:
            scores['Среднее'] = np.mean([v for k, v in scores.items() if k != 'Среднее'])
        
        return scores
    
    def evaluate_dataset(self, df: pd.DataFrame, prompt_col: str, answer_col: str) -> pd.DataFrame:
        """Evaluate all responses in the dataset"""
        evaluations = []
        eval_answers = []
        
        total_batches = (len(df) + self.config.batch_size - 1) // self.config.batch_size
        
        for i in range(0, len(df), self.config.batch_size):
            batch = df.iloc[i:i+self.config.batch_size]
            
            with tqdm(batch.iterrows(), total=len(batch), 
                     desc=f"Batch {i//self.config.batch_size + 1}/{total_batches}") as pbar:
                for _, row in pbar:
                    try:
                        scores, eval_text = self.evaluate_single_response(
                            str(row[prompt_col]), 
                            str(row[answer_col])
                        )
                        evaluations.append(scores)
                        eval_answers.append(eval_text)
                    except Exception as e:
                        print(f"Error processing row {_}: {str(e)}")
                        evaluations.append({
                            "Креативность": 0,
                            "Разнообразие": 0,
                            "Релевантность": 0,
                            "Среднее": 0
                        })
                        eval_answers.append("Error in evaluation")
                    
                    time.sleep(2)
                
            time.sleep(10)
        
        return self._create_evaluation_dataframe(df, evaluations, eval_answers)
    
    def _create_evaluation_dataframe(self, 
                                   original_df: pd.DataFrame, 
                                   evaluations: List[Dict], 
                                   eval_answers: List[str]) -> pd.DataFrame:
        score_df = pd.DataFrame(evaluations)
        df = original_df.copy()
        df['gemini_eval_answer'] = eval_answers
        return pd.concat([df, score_df], axis=1)


class StabilityEvaluator:
    def __init__(self, model_name='paraphrase-MiniLM-L6-v2'):
        self.model = SentenceTransformer(model_name)

    def calculate_similarity(self, prompts, outputs):
        prompt_embeddings = self.model.encode(prompts)
        output_embeddings = self.model.encode(outputs)
        
        similarities = cosine_similarity(prompt_embeddings, output_embeddings)
        
        stability_coefficients = np.diag(similarities)
        
        return {
            'stability_score': np.mean(stability_coefficients) * 100,  # Scale to 0-100
            'stability_std': np.std(stability_coefficients) * 100,
            'individual_similarities': stability_coefficients
        }

    def evaluate_dataset(self, df, prompt_col='rus_prompt'):
        """Evaluate stability for multiple answer columns"""
        results = {}
        
        # Find columns ending with '_answers'
        answer_columns = [col for col in df.columns if col.endswith('_answers')]
        
        for column in answer_columns:
            model_name = column.replace('_answers', '')
            results[model_name] = self.calculate_similarity(
                df[prompt_col].tolist(), 
                df[column].tolist()
            )
        
        return results


class BenchmarkEvaluator:
    def __init__(self, gemini_api_key):
        """Initialize both evaluators"""
        self.creative_evaluator = ResponseEvaluator(
            EvaluationConfig(api_key=gemini_api_key)
        )
        self.stability_evaluator = StabilityEvaluator()
        
    def evaluate_model(self, df, model_name, prompt_col='rus_prompt'):
        """Evaluate a single model's responses"""
        answer_col = f"{model_name}_answers"
        
        if answer_col not in df.columns:
            raise ValueError(f"Column {answer_col} not found in dataframe")
        
        print(f"Evaluating creativity for {model_name}...")
        creative_df = self.creative_evaluator.evaluate_dataset(df, prompt_col, answer_col)
        
        print(f"Evaluating stability for {model_name}...")
        stability_results = self.stability_evaluator.calculate_similarity(
            df[prompt_col].tolist(),
            df[answer_col].tolist()
        )
        
        creative_score = creative_df["Среднее"].mean()
        stability_score = stability_results['stability_score']
        combined_score = (creative_score + stability_score) / 2
        
        results = {
            'model': model_name,
            'creativity_score': creative_score,
            'stability_score': stability_score,
            'combined_score': combined_score,
            'creative_details': {
                'creativity': creative_df["Креативность"].mean(),
                'diversity': creative_df["Разнообразие"].mean(),
                'relevance': creative_df["Релевантность"].mean(),
            },
            'stability_details': stability_results
        }
        
        # Save detailed results
        output_file = f'evaluated_responses_{model_name}.csv'
        creative_df.to_csv(output_file, index=False)
        print(f"Detailed results saved to {output_file}")
        
        return results
    
    def evaluate_all_models(self, df, models=None, prompt_col='rus_prompt'):
        """Evaluate multiple models from the dataframe"""
        if models is None:
            # Find all columns ending with _answers
            answer_cols = [col for col in df.columns if col.endswith('_answers')]
            models = [col.replace('_answers', '') for col in answer_cols]
        
        results = []
        for model in models:
            try:
                model_results = self.evaluate_model(df, model, prompt_col)
                results.append(model_results)
                print(f"Completed evaluation for {model}")
            except Exception as e:
                print(f"Error evaluating {model}: {str(e)}")
        
        benchmark_df = pd.DataFrame(results)
        benchmark_df.to_csv('benchmark_results.csv', index=False)
        print("Benchmark completed. Results saved to benchmark_results.csv")
        
        return benchmark_df


def evaluate_single_response(gemini_api_key, prompt, response, model_name="Test Model"):
    """Evaluate a single response for the UI"""
    # Create a temporary dataframe
    df = pd.DataFrame({
        'rus_prompt': [prompt],
        f'{model_name}_answers': [response]
    })
    
    evaluator = BenchmarkEvaluator(gemini_api_key)
    
    try:
        result = evaluator.evaluate_model(df, model_name)
        
        # Format the result for displaying in UI
        output = {
            'Creativity Score': f"{result['creative_details']['creativity']:.2f}",
            'Diversity Score': f"{result['creative_details']['diversity']:.2f}",
            'Relevance Score': f"{result['creative_details']['relevance']:.2f}",
            'Average Creative Score': f"{result['creativity_score']:.2f}",
            'Stability Score': f"{result['stability_score']:.2f}",
            'Combined Score': f"{result['combined_score']:.2f}"
        }
        
        return output
    except Exception as e:
        return {
            'Error': str(e)
        }


def create_gradio_interface():
    """Create Gradio interface for evaluation app"""
    with gr.Blocks(title="Model Response Evaluator") as app:
        gr.Markdown("# Model Response Evaluator")
        gr.Markdown("Evaluate model responses for creativity, diversity, relevance, and stability.")
        
        with gr.Tab("Single Response Evaluation"):
            with gr.Row():
                gemini_api_key = gr.Textbox(label="Gemini API Key", type="password")
            
            with gr.Row():
                with gr.Column():
                    prompt = gr.Textbox(label="Original Prompt", lines=3)
                    response = gr.Textbox(label="Model Response", lines=6)
                    model_name = gr.Textbox(label="Model Name", value="Test Model")
                    
                    evaluate_btn = gr.Button("Evaluate Response")
                
                with gr.Column():
                    output = gr.JSON(label="Evaluation Results")
            
            evaluate_btn.click(
                evaluate_single_response,
                inputs=[gemini_api_key, prompt, response, model_name],
                outputs=output
            )
        
        with gr.Tab("Batch Evaluation"):
            with gr.Row():
                gemini_api_key_batch = gr.Textbox(label="Gemini API Key", type="password")
                
            with gr.Row():
                csv_file = gr.File(label="Upload CSV with responses")
                prompt_col = gr.Textbox(label="Prompt Column Name", value="rus_prompt")
                models_input = gr.Textbox(label="Model names (comma-separated, leave blank for auto-detection)")
            
            evaluate_batch_btn = gr.Button("Run Benchmark")
            benchmark_output = gr.DataFrame(label="Benchmark Results")
            
            def evaluate_batch(api_key, file, prompt_column, models_text):
                try:
                    # Load the CSV file
                    file_path = file.name
                    df = pd.read_csv(file_path)
                    
                    # Process model names if provided
                    models = None
                    if models_text.strip():
                        models = [m.strip() for m in models_text.split(',')]
                    
                    # Run the evaluation
                    evaluator = BenchmarkEvaluator(api_key)
                    results = evaluator.evaluate_all_models(df, models, prompt_column)
                    
                    return results
                except Exception as e:
                    return pd.DataFrame({'Error': [str(e)]})
            
            evaluate_batch_btn.click(
                evaluate_batch,
                inputs=[gemini_api_key_batch, csv_file, prompt_col, models_input],
                outputs=benchmark_output
            )

    return app


def main():
    parser = argparse.ArgumentParser(description="Model Response Evaluator")
    parser.add_argument("--gemini_api_key", type=str, help="Gemini API Key", default=os.environ.get("GEMINI_API_KEY"))
    parser.add_argument("--input_file", type=str, help="Input CSV file with model responses")
    parser.add_argument("--models", type=str, help="Comma-separated list of model names to evaluate")
    parser.add_argument("--prompt_col", type=str, default="rus_prompt", help="Column name containing prompts")
    parser.add_argument("--web", action="store_true", help="Launch web interface")
    
    args = parser.parse_args()
    
    if args.web:
        app = create_gradio_interface()
        app.launch(share=True)
    elif args.input_file:
        if not args.gemini_api_key:
            print("Error: Gemini API key is required. Set GEMINI_API_KEY environment variable or pass --gemini_api_key")
            return
            
        df = pd.read_csv(args.input_file)
        models = None
        if args.models:
            models = [m.strip() for m in args.models.split(',')]
            
        evaluator = BenchmarkEvaluator(args.gemini_api_key)
        evaluator.evaluate_all_models(df, models, args.prompt_col)
    else:
        print("Error: Either --input_file or --web argument is required")
        print("Run with --help for usage information")


if __name__ == "__main__":
    main()