LLM Evaluation on Google Vertex AI

gcloud components update
gcloud components install beta
gcloud auth login

pip install google-cloud google-cloud-aiplatform openai wandb weave arxiv pymupdf rouge-score

import os
import arxiv
import fitz  # PyMuPDF
import json
from vertexai.generative_models import GenerativeModel, GenerationConfig
import vertexai
import weave; weave.init('paper_abstract_gen')
import re
from time import sleep


# Set up Vertex AI
PROJECT_ID = "dsports-6ab79"
LOCATION = "us-central1"
vertexai.init(project=PROJECT_ID, location=LOCATION)


# Directory to save downloaded papers
download_dir = "arxiv_papers"
os.makedirs(download_dir, exist_ok=True)


# Define AI-specific search queries
search_queries = [
    "Large Language Models for vision tasks AND cat:cs.AI",
    "Multimodal AI techniques AND cat:cs.CV",
    "Applications of Transformers in healthcare AI AND cat:cs.LG",
    "Few-shot learning in AI and ML AND cat:cs.LG",
    "Vision and language models integration AND cat:cs.CV",
    "Domain-specific fine-tuning for ML models AND cat:cs.LG",
    "Foundational models in AI and CV applications AND cat:cs.AI",
    "NLP in robotics and vision systems AND cat:cs.AI",
    "Bias and fairness in AI for CV AND cat:cs.CV",
    "Evaluation metrics for multimodal AI AND cat:cs.LG"
]


def download_papers(max_pages=15, max_attempts_per_query=20):
    """Download papers or use existing ones from the download directory."""
    papers = []
    downloaded_titles = set()
    
    # First check for existing papers
    if os.path.exists(download_dir):
        existing_pdfs = [f for f in os.listdir(download_dir) if f.endswith('.pdf')]
        if existing_pdfs:
            print(f"Found {len(existing_pdfs)} existing papers, checking validity...")
            for pdf_file in existing_pdfs:
                pdf_path = os.path.join(download_dir, pdf_file)
                try:
                    with fitz.open(pdf_path) as pdf:
                        if pdf.page_count <= max_pages:
                            arxiv_id = pdf_file.replace('.pdf', '')
                            # Try to get a clean title from the first page
                            title = pdf[0].get_text().split('\n')[0].strip()
                            papers.append({
                                "title": title,
                                "file_path": pdf_path,
                                "arxiv_id": arxiv_id
                            })
                            downloaded_titles.add(title)
                            print(f"Using existing paper: {title}")
                except Exception as e:
                    print(f"Error checking existing PDF {pdf_path}: {e}")
                    if os.path.exists(pdf_path):
                        os.remove(pdf_path)
    
    # If we have enough papers (one per query), return early
    if len(papers) >= len(search_queries):
        print(f"\nUsing {len(papers)} existing papers")
        return papers[:len(search_queries)]  # Return only what we need
    
    # Otherwise, download remaining papers
    print(f"\nNeed {len(search_queries) - len(papers)} more papers, downloading...")
    
    client = arxiv.Client()
    
    for query in search_queries[len(papers):]:  # Only process remaining queries
        paper_found = False
        attempt = 0
        
        while not paper_found and attempt < max_attempts_per_query:
            search = arxiv.Search(
                query=query,
                max_results=100,
                sort_by=arxiv.SortCriterion.SubmittedDate
            )
            
            try:
                results = list(client.results(search))
                start_idx = attempt * 5
                end_idx = start_idx + 5
                current_batch = results[start_idx:end_idx]
                
                for result in current_batch:
                    if result.title not in downloaded_titles:
                        print(f"Downloading: {result.title}")
                        paper_id = result.entry_id.split('/')[-1]
                        pdf_filename = f"{paper_id}.pdf"
                        pdf_path = os.path.join(download_dir, pdf_filename)
                        
                        result.download_pdf(dirpath=download_dir, filename=pdf_filename)
                        
                        try:
                            with fitz.open(pdf_path) as pdf:
                                if pdf.page_count <= max_pages:
                                    papers.append({
                                        "title": result.title,
                                        "file_path": pdf_path,
                                        "arxiv_id": paper_id
                                    })
                                    downloaded_titles.add(result.title)
                                    print(f"Accepted: {result.title}")
                                    paper_found = True
                                    break
                                else:
                                    os.remove(pdf_path)
                                    print(f"Skipped (too many pages: {pdf.page_count}): {result.title}")
                        except Exception as e:
                            print(f"Error checking PDF {pdf_path}: {e}")
                            if os.path.exists(pdf_path):
                                os.remove(pdf_path)
                
                attempt += 1
                if not paper_found:
                    print(f"Attempt {attempt}/{max_attempts_per_query} for query: {query}")
                    sleep(3)
                    
            except Exception as e:
                print(f"Error during download: {e}")
                sleep(3)
                attempt += 1
                continue
        
        if not paper_found:
            print(f"Failed to find suitable paper for query after {max_attempts_per_query} attempts: {query}")
    
    print(f"\nTotal papers available: {len(papers)} ({len(papers) - len(search_queries)} existing, {len(search_queries) - (len(papers) - len(search_queries))} new)")
    return papers


def extract_first_page_text(pdf_path):
    """Extract text from only the first page of the PDF."""
    with fitz.open(pdf_path) as pdf:
        if pdf.page_count > 0:
            page = pdf[0]
            return page.get_text()
    return ""

def extract_abstract_with_gemini(text, title, max_retries=3):
    """Extract abstract using Gemini with simple retry logic."""
    model = GenerativeModel("gemini-1.5-pro-002")
    
    prompt = (
        f"From the following first page of the research paper titled '{title}', "
        f"extract ONLY the abstract section. Return the result in JSON format with 'abstract' as the key. "
        f"If you cannot find the abstract, return an empty string as the value.\n\n"
        f"Paper content:\n\n{text}"
    )
    
    for attempt in range(max_retries):
        try:
            response = model.generate_content(
                prompt,
                generation_config=GenerationConfig(
                    temperature=0,
                    response_mime_type="application/json"
                )
            )
            return json.loads(response.text)
            
        except Exception as e:
            if attempt == max_retries - 1:  # Last attempt
                print(f"Error extracting abstract for {title}: {e}")
                return {"abstract": ""}
            
            # Simple exponential backoff: 5s, 10s, 20s
            sleep_time = 5 * (2 ** attempt)
            print(f"Attempt {attempt + 1} failed, retrying in {sleep_time}s...")
            sleep(sleep_time)
    
    return {"abstract": ""}

def count_words(text):
    """Count words excluding punctuation and special characters."""
    cleaned_text = re.sub(r'[^\w\s]', ' ', text.lower())
    words = [word for word in cleaned_text.split() if word.strip()]
    return len(words)


def main():
    # Download papers
    papers = download_papers()
    print(f"\nDownloaded {len(papers)} papers. Processing abstracts...\n")
    
    # Process papers and extract abstracts
    paper_data = []
    for paper in papers:
        title = paper["title"]
        pdf_path = paper["file_path"]
        
        print(f"Processing: {title}")
        first_page_text = extract_first_page_text(pdf_path)
        abstract_json = extract_abstract_with_gemini(first_page_text, title)
        abstract_text = abstract_json.get('abstract', '')
        word_count = count_words(abstract_text)
        
        paper_data.append({
            "title": title,
            "file_path": pdf_path,
            "abstract": abstract_text,
            "word_count": word_count,
            "arxiv_id": paper["arxiv_id"]
        })
        
        sleep(2)


    # Save to JSONL file
    output_file = "paper_abstracts.jsonl"
    with open(output_file, "w") as f:
        for entry in paper_data:
            json.dump(entry, f)
            f.write("\n")


    print(f"\nProcessed {len(paper_data)} papers. Results saved to {output_file}")


if __name__ == "__main__":
    main()


import json
import fitz  # PyMuPDF
import time
from vertexai.generative_models import GenerativeModel, GenerationConfig
import vertexai
import weave; weave.init('vertex_abstract_prediction')


# Configuration
PROJECT_ID = "dsports-6ab79"
LOCATION = "us-central1"
vertexai.init(project=PROJECT_ID, location=LOCATION)


# Model configurations
MODELS = {
    "gemini_2_flash": {
        "name": "gemini-2.0-flash-001",
        "type": "vertex",
        "delay": 2,
        "temperature": 0.0
    },
    "gemini_1_5_flash": {
        "name": "gemini-1.5-flash-002",
        "type": "vertex",
        "delay": 2,
        "temperature": 0.0
    }
}


def load_paper_data():
    """Load papers with their abstracts and word counts."""
    with open("paper_abstracts.jsonl", "r") as f:
        return [json.loads(line) for line in f]


def extract_text_after_page_one(pdf_path):
    """Extract text from page 2 onwards."""
    text = ""
    try:
        with fitz.open(pdf_path) as pdf:
            if pdf.page_count > 1:
                for page_num in range(1, pdf.page_count):  # Start from page 2
                    page = pdf[page_num]
                    text += page.get_text()
    except Exception as e:
        print(f"Error extracting text from PDF {pdf_path}: {e}")
    return text


def create_abstract_prompt(text, title, target_length):
    """Create a prompt for abstract generation with target length."""
    return (
        f"You are tasked with generating an abstract for a research paper titled '{title}'. "
        f"The abstract should be approximately {target_length} words long.\n\n"
        f"Generate an abstract that summarizes the key points of the paper, including the "
        f"research objective, methodology, and main findings. The abstract should be "
        f"self-contained and clearly communicate the paper's contribution. Respond only with the ABSTRACT, and NOT the title\n\n"
        f"Paper content:\n\n{text}"
        f"Respond only with the ABSTRACT!"
    )


@weave.op
def predict_abstract_with_gemini(text, title, target_length, model_info):
    """Generate abstract prediction using specified Gemini model."""
    model = GenerativeModel(model_info["name"])
    generation_config = GenerationConfig(
        temperature=model_info["temperature"]
    )
    
    try:
        response = model.generate_content(
            create_abstract_prompt(text, title, target_length),
            generation_config=generation_config
        )
        time.sleep(model_info["delay"])
        return response.text
    except Exception as e:
        print(f"Error generating abstract with {model_info['name']}: {e}")
        return ""


def process_papers(start_index=0, end_index=None):
    """Process papers and generate abstract predictions."""
    papers = load_paper_data()
    model_predictions = {model_name: [] for model_name in MODELS.keys()}
    
    papers_to_process = papers[start_index:end_index] if end_index else papers[start_index:]
    
    for i, paper in enumerate(papers_to_process, start=start_index):
        title = paper["title"]
        pdf_path = paper["file_path"]
        target_length = paper["word_count"]
        
        print(f"\nProcessing paper {i+1}/{len(papers)}: {title}")


        # Extract text from page 2 onwards
        paper_text = extract_text_after_page_one(pdf_path)
        if not paper_text:
            print(f"Skipping file {pdf_path} - no text found after page 1.")
            continue


        # Generate predictions using each model
        for model_name, model_info in MODELS.items():
            try:
                print(f"Generating abstract prediction using {model_name}...")
                predicted_abstract = predict_abstract_with_gemini(
                    paper_text, 
                    title, 
                    target_length, 
                    model_info
                )
                
                model_predictions[model_name].append({
                    "title": title,
                    "file_path": pdf_path,
                    "abstract": predicted_abstract
                })
                
                print(f"Successfully generated {model_name} abstract prediction")
                
            except Exception as e:
                print(f"Error processing paper {title} with {model_name}: {e}")


        # Save progress after each paper
        for model_name, predictions in model_predictions.items():
            output_file = f"{model_name}_abstract_predictions.jsonl"
            with open(output_file, "w") as f:
                for entry in predictions:
                    json.dump(entry, f)
                    f.write("\n")


    return model_predictions


if __name__ == "__main__":
    print("Starting abstract prediction...")
    print(f"Using temperature settings:")
    for model, config in MODELS.items():
        print(f"- {model}: temperature={config['temperature']}, delay={config['delay']}s")
    
    predictions = process_papers()
    print("\nProcessing completed. Predictions saved to individual JSONL files.")


import weave
from weave import Model
import json
import vertexai
from vertexai.generative_models import GenerativeModel, GenerationConfig
from time import sleep
import asyncio
from rouge_score.rouge_scorer import RougeScorer
from typing import Dict, Any
import bert_score


# Initialize Vertex AI and Weave
PROJECT_ID = "dsports-6ab79"
LOCATION = "us-central1"
vertexai.init(project=PROJECT_ID, location=LOCATION)
weave.init('abstract_metrics_eval')


class BaseJsonModel(Model):
    """Base model class for loading abstracts from JSON files."""
    abstract_file: str = ""


    def get_abstracts(self) -> dict:
        """Load abstracts from the JSON file."""
        abstracts = {}
        with open(self.abstract_file, 'r') as f:
            for line in f:
                entry = json.loads(line)
                abstracts[entry['title']] = entry['abstract']
        return abstracts


class GeminiFlash_1_5_Model(BaseJsonModel):
    """Specific model class for Gemini Flash 1.5 abstracts."""
    @weave.op
    def predict(self, title: str) -> dict:
        """Return the pre-generated abstract for a given title."""
        abstracts = self.get_abstracts()
        return {"model_output": abstracts.get(title, "")}


class GeminiFlash2Model(BaseJsonModel):
    """Specific model class for Gemini Flash abstracts."""


    @weave.op
    def predict(self, title: str) -> dict:
        """Return the pre-generated abstract for a given title."""
        abstracts = self.get_abstracts()
        return {"model_output": abstracts.get(title, "")}


@weave.op
def bert_scorer(gt_abstract: str, model_output: dict) -> Dict[str, float]:
    """Calculate BERTScore for the abstract."""
    if not model_output or 'model_output' not in model_output:
        return {'bert_score': 0.0}
    
    try:
        P, R, F1 = bert_score.score(
            [model_output['model_output']],
            [gt_abstract],
            lang='en',
            model_type='microsoft/deberta-xlarge-mnli'
        )
        return {'bert_score': float(F1.mean())}
    except Exception as e:
        print(f"Error calculating BERTScore: {e}")
        return {'bert_score': 0.0}


async def gemini_scorer(gt_abstract: str, model_output: dict) -> dict:
    """Evaluate abstract using Gemini model."""
    if not model_output or 'model_output' not in model_output:
        print("Invalid model output")
        return {'gemini_score': 0}

    model_id = "gemini-1.5-pro-002"
    model = GenerativeModel(model_id)
    
    response_schema = {
        "type": "object",
        "properties": {
            "score": {"type": "integer", "minimum": 1, "maximum": 5}
        },
        "required": ["score"]
    }

    formatted_text = (
        f"Given these two research paper abstracts:\n\n"
        f"Ground Truth Abstract:\n{gt_abstract}\n\n"
        f"Generated Abstract:\n{model_output['model_output']}\n\n"
        f"Rate how well the generated abstract captures the key information from the ground truth abstract "
        f"on a scale from 1-5, where 1 is poor and 5 is excellent. Consider:\n"
        f"Respond with ONLY a JSON object in this format: {{'score': X}} where X is your integer rating."
    )
    
    max_attempts = 10
    attempt = 0
    base_delay = 3
    
    while attempt < max_attempts:
        print(f"evaluating with gemini attempt: {attempt}")
        try:
            response = model.generate_content(
                formatted_text,
                generation_config=GenerationConfig(
                    temperature=0.0,
                    response_mime_type="application/json",
                    response_schema=response_schema
                )
            )
            
            eval_result = json.loads(response.text)
            score = int(eval_result.get('score', 0))
            
            if not 1 <= score <= 5:
                raise ValueError(f"Invalid score: {score}")
                
            print(f"Sleeping for {base_delay}s between calls")
            sleep(base_delay)
            return {'gemini_score': score}
            
        except Exception as e:
            attempt += 1
            if "429" in str(e) and attempt < max_attempts:
                delay = base_delay * (2 ** attempt)
                print(f"Rate limit hit. Attempt {attempt}/{max_attempts}. "
                      f"Retrying in {delay}s...")
                sleep(delay)
            else:
                if attempt == max_attempts:
                    print(f"Max attempts reached")
                print(f"Error in evaluation: {e}")
                return {'gemini_score': 0}
    
    return {'gemini_score': 0}

@weave.op
def rouge_scorer(gt_abstract: str, model_output: dict) -> Dict[str, float]:
    """Calculate ROUGE scores for the abstract."""
    if not model_output or 'model_output' not in model_output:
        return {
            'rouge1_f': 0.0,
            'rouge2_f': 0.0,
            'rougeL_f': 0.0
        }
    
    try:
        scorer = RougeScorer(['rouge1', 'rouge2', 'rougeL'], use_stemmer=True)
        scores = scorer.score(gt_abstract, model_output['model_output'])
        
        return {
            'rouge1_f': float(scores['rouge1'].fmeasure),
            'rouge2_f': float(scores['rouge2'].fmeasure),
            'rougeL_f': float(scores['rougeL'].fmeasure)
        }
    except Exception as e:
        print(f"Error calculating ROUGE scores: {e}")
        return {
            'rouge1_f': 0.0,
            'rouge2_f': 0.0,
            'rougeL_f': 0.0
        }


@weave.op
def compression_scorer(gt_abstract: str, model_output: dict) -> Dict[str, float]:
    """Calculate compression ratio of the abstract."""
    if not model_output or 'model_output' not in model_output:
        return {'compression_ratio': 0.0}
    
    try:
        gt_words = len(gt_abstract.split())
        generated_words = len(model_output['model_output'].split())
        
        compression_ratio = min(gt_words, generated_words) / max(gt_words, generated_words)
        
        return {'compression_ratio': float(compression_ratio)}
    except Exception as e:
        print(f"Error calculating compression ratio: {e}")
        return {'compression_ratio': 0.0}


@weave.op
def coverage_scorer(gt_abstract: str, model_output: dict) -> Dict[str, float]:
    """Calculate content coverage using word overlap."""
    if not model_output or 'model_output' not in model_output:
        return {'coverage_score': 0.0}
    
    try:
        gt_words = set(gt_abstract.lower().split())
        generated_words = set(model_output['model_output'].lower().split())
        
        intersection = len(gt_words.intersection(generated_words))
        union = len(gt_words.union(generated_words))
        
        coverage_score = intersection / union if union > 0 else 0.0
        
        return {'coverage_score': float(coverage_score)}
    except Exception as e:
        print(f"Error calculating coverage score: {e}")
        return {'coverage_score': 0.0}


def create_evaluation_dataset(gt_file: str):
    """Create dataset from ground truth file."""
    dataset = []
    with open(gt_file, 'r') as f:
        for line in f:
            entry = json.loads(line)
            dataset.append({
                "title": entry["title"],
                "gt_abstract": entry["abstract"]
            })
    return dataset


async def run_evaluations(gt_file: str):
    """Run separate evaluations for each model."""
    eval_dataset = create_evaluation_dataset(gt_file)
    scorers = [
        gemini_scorer,
        rouge_scorer,
        compression_scorer,
        coverage_scorer,
        bert_scorer
    ]
    
    # Create and evaluate Gemini flash 1.5 model
    print("\nEvaluating Gemini Flash 1.5 abstracts...")
    flash_1_5_model = GeminiFlash_1_5_Model(abstract_file="gemini_1_5_flash_abstract_predictions.jsonl")
    flash_1_5_evaluation = weave.Evaluation(
        dataset=eval_dataset,
        scorers=scorers
    )
    flash1_5_results = await flash_1_5_evaluation.evaluate(flash_1_5_model)
    
    # Create and evaluate Gemini Flash model
    print("\nEvaluating Gemini Flash abstracts...")
    flash_model = GeminiFlash2Model(abstract_file="gemini_2_flash_abstract_predictions.jsonl")
    flash_2_evaluation = weave.Evaluation(
        dataset=eval_dataset,
        scorers=scorers
    )
    flash2_results = await flash_2_evaluation.evaluate(flash_model)
    
    # Print results
    print("\nEvaluation Results:")
    print("\nGemini Flash 1.5 Results:")
    print(json.dumps(flash1_5_results, indent=2))
    print("\nGemini Flash 2.0 Results:")
    print(json.dumps(flash2_results, indent=2))
    
    return {
        "gemini_flash1_5": flash1_5_results,
        "gemini_flash2": flash2_results
    }


if __name__ == "__main__":
    gt_file = "paper_abstracts.jsonl"
    asyncio.run(run_evaluations(gt_file))