from openai import OpenAI import os import gzip import json from collections import Counter import re import threading from concurrent.futures import ThreadPoolExecutor, as_completed import time from tqdm import tqdm from typing import Literal, List, Dict, Any _global_config = { 'api_key': '0', 'base_url': 'http://localhost:8000/v1', 'model': None, 'iteration': 10, 'instruction_template': ''' You are a TypeScript type inference expert. Please perform type inference for the part based on the code slice below, and output the result strictly in the required format. ''' } _thread_local = threading.local() def read_jsonl_gz(filename): data = [] with gzip.open(filename, 'rt', encoding='utf-8') as f: for line in f: data.append(json.loads(line)) return data def process_single_item(args): item, idx, total_items = args item_id = item.get('id', '') if not hasattr(_thread_local, "client"): _thread_local.client = OpenAI(api_key=_global_config['api_key'], base_url=_global_config['base_url']) client = _thread_local.client model = _global_config['model'] iteration = _global_config['iteration'] instruction_template = _global_config['instruction_template'] responses = [] for i in range(iteration): context = f"{instruction_template}\nCode slice: {item['sliced_code']}" try: messages = [{"role": "user", "content": context}] result = client.chat.completions.create( messages=messages, model=model, temperature=0.5, top_p=0.9, max_tokens=128 ) resp_str = result.choices[0].message.content.strip() resp_str = resp_str.replace('\n', '') except Exception as e: print(f"API call error: {e}") return None match = re.search(r'\s*([^\n\r]+)', resp_str) if match: extracted_type = f"{match.group(1).strip()}" else: extracted_type = resp_str responses.append(extracted_type) response_counter = Counter(responses) predictions = [] for pred_type, count in response_counter.most_common(): frequency = count / iteration predictions.append([pred_type, frequency]) result = { 'id': item_id, 'predictions': predictions } return result def save_results_thread_safe(results, output_fp, lock): with lock: try: with open(output_fp, 'w', encoding='utf-8') as f: json.dump(results, f, ensure_ascii=False, indent=4) except Exception as e: print(f"Error writing file: {e}") def append_results_jsonl_thread_safe(buffer: List[Dict[str, Any]], output_fp: str, lock: threading.Lock): if not buffer: return with lock: try: with open(output_fp, 'a', encoding='utf-8') as f: for item in buffer: f.write(json.dumps(item, ensure_ascii=False) + '\n') except Exception as e: print(f"Error writing file: {e}") def prediction(testdata_fp, model, iteration=10, output_fp=None, num_threads=4, output_format: Literal['json', 'jsonl', None] = None): _global_config['model'] = model _global_config['iteration'] = iteration if output_fp is not None: output_dir = os.path.dirname(output_fp) if output_dir and not os.path.exists(output_dir): os.makedirs(output_dir) if output_format is None: if output_fp and output_fp.endswith('.jsonl'): output_format = 'jsonl' else: output_format = 'json' existing_ids = set() existing_results = [] if output_fp is not None and os.path.exists(output_fp): try: if output_format == 'jsonl': with open(output_fp, 'r', encoding='utf-8') as f: for line in f: if not line.strip(): continue obj = json.loads(line) existing_results.append(obj) if 'id' in obj: existing_ids.add(obj['id']) else: with open(output_fp, 'r', encoding='utf-8') as f: existing_results = json.load(f) existing_ids = {item.get('id', '') for item in existing_results if 'id' in item} print(f"Loaded {len(existing_ids)} existing results") except (json.JSONDecodeError, FileNotFoundError) as e: print(f"Error reading existing file: {e}, creating new file") existing_results = [] existing_ids = set() data = read_jsonl_gz(testdata_fp) items_to_process = [] for idx, item in enumerate(data): item_id = item.get('id', '') if item_id not in existing_ids: items_to_process.append((item, idx, len(data))) else: print(f"[{idx+1}/{len(data)}] Skipping completed id: {item_id}") if not items_to_process: print("All items have been processed!") return existing_results print(f"Processing {len(items_to_process)} items with {num_threads} threads") results = existing_results.copy() results_lock = threading.Lock() results_list = results pending_buffer: List[Dict[str, Any]] = [] start_time = time.time() skipped_due_error = 0 last_save_time = start_time with ThreadPoolExecutor(max_workers=num_threads) as executor: future_to_item = {executor.submit(process_single_item, item): item for item in items_to_process} with tqdm(total=len(items_to_process), desc="Progress", unit="items") as pbar: completed_count = 0 for future in as_completed(future_to_item): try: result = future.result() if result is None: skipped_due_error += 1 else: with results_lock: results_list.append(result) if output_format == 'jsonl': pending_buffer.append(result) completed_count += 1 pbar.update(1) if output_fp and (completed_count % 100 == 0 or time.time() - last_save_time > 120): if output_format == 'jsonl': with results_lock: buffer_to_flush = list(pending_buffer) pending_buffer.clear() append_results_jsonl_thread_safe(buffer_to_flush, output_fp, results_lock) else: with results_lock: save_results = list(results_list) save_results_thread_safe(save_results, output_fp, results_lock) pbar.set_postfix({"saved": "checkpoint"}) last_save_time = time.time() except Exception as e: print(f"Error processing task: {e}") final_results = results_list if output_fp is not None: if output_format == 'jsonl': with results_lock: buffer_to_flush = list(pending_buffer) pending_buffer.clear() append_results_jsonl_thread_safe(buffer_to_flush, output_fp, results_lock) else: save_results_thread_safe(final_results, output_fp, results_lock) skipped_count = len(data) - len(items_to_process) print(f"\nDone! Total: {len(data)}, skipped: {skipped_count}, new: {len(final_results) - len(existing_results)}, errors: {skipped_due_error}") return final_results if __name__ == "__main__": testdata_fp = './sliced_data_testdata.jsonl.gz' model = 'my_qwen3-8b' iteration = 20 output_fp = './result_my/testdata_prediction.jsonl' num_threads = 32 prediction( testdata_fp=testdata_fp, model=model, iteration=iteration, output_fp=output_fp, num_threads=num_threads )