ruby56/Citation-Benchmark

1import sqlite32import random3 4def build_tasks():5    print("Generating Academic Peer Reviewer Tasks...")6    try:7        conn = sqlite3.connect('citation_db.sqlite')8        cursor = conn.cursor()9        10        # Pull 50 random papers that have AT LEAST ONE citation in our db11        query = """12        SELECT s.corpus_id, a.title, a.abstract13        FROM s2_papers s14        JOIN arxiv_metadata a ON s.arxiv_id = a.arxiv_id15        WHERE EXISTS (SELECT 1 FROM citations c WHERE c.cited_corpus_id = s.corpus_id)16        LIMIT 5017        """18        cursor.execute(query)19        papers = cursor.fetchall()20        21        if len(papers) < 50:22             print("WARNING: Less than 50 cited papers found, utilizing available subset.")23             24        # Grab 25 random abstracts to use as decoys25        cursor.execute("SELECT abstract FROM arxiv_metadata ORDER BY RANDOM() LIMIT 25")26        decoy_abstracts = [r[0] for r in cursor.fetchall()]27 28        tasks_str = "import json\nfrom pydantic import BaseModel\nfrom typing import List\n\n"29        tasks_str += "class Task(BaseModel):\n    id: str\n    target_paper_id: str\n    target_paper_title: str\n    target_abstract: str\n    ground_truth_rating: str\n\n"30        tasks_str += "TASKS = [\n"31        32        for i, paper in enumerate(papers):33            corpus_id, title, actual_abstract = paper34            title = title.replace('"', "'").replace('\n', ' ') if title else ""35            actual_abstract = actual_abstract.replace('"', "'").replace('\n', ' ') if actual_abstract else ""36            37            task_id = f"T{i+1:03d}"38            39            if i < 25:40                # Valid Paper41                rating = "ACCEPT"42                abstract = actual_abstract43            else:44                # Decoy Paper (Abstract is swapped with a random one)45                rating = "REJECT"46                abstract = decoy_abstracts[i - 25].replace('"', "'").replace('\n', ' ') if decoy_abstracts[i-25] else "Fake Abstract."47                48            tasks_str += f"""    Task(49        id="{task_id}",50        target_paper_id="{corpus_id}",51        target_paper_title="{title}",52        target_abstract="{abstract}",53        ground_truth_rating="{rating}"54    ),\n"""55 56        tasks_str += "]\n\n"57        58        # Grader Logic for Accept/Reject59        tasks_str += """class Grader:60    def __init__(self, task: Task):61        self.task = task62        63    def score(self, rating: str, db_conn=None) -> float:64        if rating.strip().upper() == self.task.ground_truth_rating.strip().upper():65            return 1.066        return 0.067"""68        69        with open('tasks.py', 'w', encoding='utf-8') as f:70            f.write(tasks_str)71            72        print(f"Successfully generated {len(papers)} Peer Review tasks into tasks.py!")73        conn.close()74        75    except Exception as e:76        print(f"Failed: {e}")77 78if __name__ == "__main__":79    build_tasks()80