File size: 6,843 Bytes
e76659b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
from __future__ import annotations
import pandas as pd
from openai import OpenAI
from concurrent.futures import ThreadPoolExecutor, as_completed
from time import sleep
from typing import List, Dict, Any, Optional
from openai import OpenAI
import json
from pathlib import Path
from typing import Any, Dict, Iterable, List, Union
import re
import datetime
from typing import Dict, List, Optional
from mathruler.grader import extract_boxed_content, grade_answer
import math
from tqdm.auto import tqdm


def extract_description(predict: str) -> Optional[str]:
    """
    Extracts the content of the <answer>…</answer> block from `predict`.
    Returns the inner text (with leading/trailing whitespace stripped),
    or None if no <answer> tag is found.
    """
    match = re.search(r"<description>([\s\S]*?)</description>", predict, re.DOTALL)
    if not match:
        return predict
    return match.group(1).strip()

# curl http://29.81.228.243:8081 /v1/models
client = OpenAI(
    base_url="http://29.81.244.54:8081/v1",  # your vLLM server
    api_key="ANYKEY",                        # if you set --api-key when launching
)


def chat_once(messages):
    resp = client.chat.completions.create(
        model="Qwen2.5-VL-72B-Instruct",
        messages=messages
    )
    return resp.choices[0].message.content


def chat_batch(
    client,
    all_message_batches: List[List[Dict[str, str]]],
    *,
    model: str = "Qwen2.5-VL-72B-Instruct",
    max_workers: int = 8,
    retries: int = 2,
    backoff: float = 0.5,
    timeout: Optional[float] = None,
) -> List[str]:
    """
    Send many chat requests in parallel and return replies as a list of strings,
    preserving the order of `all_message_batches`.
    """

    def _chat_once_with_retry(messages: List[Dict[str, str]]) -> str:
        last_err: Optional[BaseException] = None
        for attempt in range(retries + 1):
            try:
                resp = client.chat.completions.create(
                    model=model,
                    messages=messages,
                    timeout=timeout,
                )
                # Different SDKs expose content slightly differently; handle common cases.
                choice = resp.choices[0]
                if hasattr(choice, "message") and getattr(choice.message, "content", None) is not None:
                    return choice.message.content
                if hasattr(choice, "text") and choice.text is not None:
                    return choice.text
                # Fallback to stringifying the choice if structure is unexpected.
                return str(choice)
            except Exception as e:
                last_err = e
                if attempt < retries:
                    sleep(backoff * (2 ** attempt))
        return f"Error: {last_err!r}"

    results: List[Optional[str]] = [None] * len(all_message_batches)
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_idx = {
            executor.submit(_chat_once_with_retry, batch): i
            for i, batch in enumerate(all_message_batches)
        }
        for fut in as_completed(future_to_idx):
            i = future_to_idx[fut]
            results[i] = fut.result()

    # mypy-friendly cast: no Nones remain at this point
    return [r if r is not None else "Error: Unknown failure" for r in results]



def load_json_list(path: Union[str, Path], encoding: str = "utf-8") -> List[Dict[str, Any]]:
    """
    Load a JSON file whose top-level structure is a list of dicts.

    Raises:
        FileNotFoundError, json.JSONDecodeError, TypeError
    """
    p = Path(path)
    with p.open("r", encoding=encoding) as f:
        data = json.load(f)

    if not isinstance(data, list):
        raise TypeError(f"Expected top-level JSON to be a list, got {type(data).__name__}")

    for i, item in enumerate(data):
        if not isinstance(item, dict):
            raise TypeError(f"Item at index {i} is {type(item).__name__}, expected dict")

    return data

# Prepare a list of different message‐lists you want to send:
all_message_batches = [
    [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user",   "content": "Hello, how are you?"}
    ],
    [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user",   "content": "Tell me a joke."}
    ],
    [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user",   "content": "Tell me a joke."}
    ],
    [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user",   "content": "Tell me a joke."}
    ],
    # …more batches…
]


res = chat_batch(client, all_message_batches)

prompt_template = '''Text description: {Description}\nQuestion: {Question}\nYou are provided a text description of a problem and a question. Determine the answer to the question based on the text description. First provide an internal step-by-step reasoning within <think> </think> tags, then provide a single word or phrase answer in \\boxed{}.'''
MODEL = "Qwen2.5-VL-72B-Instruct"
BATCH_SIZE = 16
filename = "MLLM_rlvr_train"
out_file = f'./caption_out/{filename}.json'
data = load_json_list(f'./gemini-flash/{filename}.json')


def to_messages(example: Dict[str, Any]) -> List[Dict[str, str]]:
    """Use the single string inside `predictions` as the user input."""
    preds = example.get("predictions")
    question = example.get("problem")
    
    if isinstance(preds, list) and preds:
        first = preds[0]
        text = first if isinstance(first, str) else json.dumps(first, ensure_ascii=False)
        description = extract_description(text)
        input_question = prompt_template.replace('{Description}', description).replace('{Question}', question)
    else:
        input_question = 'None'

    return [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user",   "content": input_question},
    ]


# Ensure output dir exists and start fresh
Path(out_file).parent.mkdir(parents=True, exist_ok=True)
with open(out_file, "w", encoding="utf-8"):
    pass

total = len(data)
num_batches = math.ceil(total / BATCH_SIZE)

for start in tqdm(range(0, total, BATCH_SIZE),
                  total=num_batches, desc="Batches", unit="batch"):
    chunk = data[start : start + BATCH_SIZE]
    batch_messages = [to_messages(ex) for ex in chunk]

    replies = chat_batch(client, batch_messages, model=MODEL,
                         max_workers=8, retries=2, backoff=0.5, timeout=None)

    print(replies[0])
    with open(out_file, "a", encoding="utf-8") as f:
        for ex, reply in zip(chunk, replies):
            record = {**ex, "model": MODEL, "model_caption_response": reply}
            f.write(json.dumps(record, ensure_ascii=False) + "\n")
        f.flush()