score.py

import json
import jsonlines
from openai import OpenAI
import concurrent.futures
import threading
from tqdm import tqdm
import os
import argparse

class Counter:
    def __init__(self):
        self.hit = 0
        self.total = 0
        self.lock = threading.Lock()
    
    def increment_hit(self):
        with self.lock:
            self.hit += 1
    
    def increment_total(self):
        with self.lock:
            self.total += 1

def process_item(item, index, client, prompt, writer, counter):
    try:
        answer = item["answer"]
        gt = item['gt']            
        text = f"Answer: {answer} Groundtruth: {gt}"

        conversation_history = [
            {"role": "system",
             "content": [{"type": "text", "text": prompt}]},
            {"role": "user",
             "content": [{"type": "text", "text": text}]}
        ]
        
        response = client.chat.completions.create(
            model='gpt-4o-mini-2024-07-18',
            messages=conversation_history,
            max_tokens=4096,
            temperature=0
        )
        answer_corrected = response.choices[0].message.content.strip()
        
        if answer_corrected not in ["0", "1"]:
            answer_corrected = "0" 
        
        with threading.Lock():
            writer.write({
                'number': index, 
                "answer": answer, 
                "gt": gt, 
                'hit': answer_corrected
            })
        
        counter.increment_total()
        if answer_corrected == "1":
            counter.increment_hit()
            
    except Exception as e:
        print(f"Error in {index}: {str(e)}")
        with threading.Lock():
            writer.write({
                'number': index, 
                "answer": answer, 
                "gt": gt, 
                'hit': f"error: {str(e)}"
            })
        counter.increment_total()

def main(args):
    api_keys = os.getenv("OPENAI_API_KEY")
    base_url = os.getenv("OPENAI_API_BASE")

    prompt = "Compare the given two two-dimensional matrices. Allow for differences in format, but each value should be identical. If the two matrices are identical, reply with 1; otherwise, reply with 0."
    input_file = args.input_file
    output_file = args.output_file
    max_workers = args.num_workers 
    
    client = OpenAI(
        api_key=api_keys, 
        base_url=base_url
    )
    
    with open(input_file, 'r', encoding='utf-8') as f:
        items = list(jsonlines.Reader(f))
    
    counter = Counter()
    
    with jsonlines.open(output_file, mode='w') as writer, \
        concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        
        pbar = tqdm(total=len(items), desc="Processing")
        
        futures = []
        for i, item in enumerate(items):
            future = executor.submit(
                process_item, 
                item, i, client, prompt, writer, counter
            )
            future.add_done_callback(lambda _: pbar.update(1))
            futures.append(future)
        
        for future in concurrent.futures.as_completed(futures):
            pass  
        
        pbar.close()
    
    print(f"Hit: {counter.hit}")
    print(f"Total: {counter.total}")
    print(f"Acc: {counter.hit / counter.total:.4f}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Define hyper-parameters for reasoning')
    parser.add_argument('--input_file', default="result.json", type=str, help='results after inference')
    parser.add_argument('--output_file', default="result_score.json", type=str, help='scored results after scoring')
    parser.add_argument('--num_workers', default=40, type=int, help='num workers in scoring')
    args = parser.parse_args()
    main(args)