image_generation.py

import os, uuid, subprocess, sys, textwrap, re
import base64, json, ast, resource

os.environ["HF_HOME"] = "/export/projects/nlp/.cache"
os.environ["HF_HUB_DISABLE_XET"] = "1"

from vllm import LLM, SamplingParams
from vllm.sampling_params import StructuredOutputsParams
from pydantic import BaseModel, Field
import matplotlib.image as image

MODEL_NAME = "Qwen/Qwen3.6-27B-FP8"

model = None

def load_model():
    global model
    model = LLM(model=MODEL_NAME, max_model_len=8192, max_num_seqs=256)

class ReviewOutput(BaseModel):
	reasoning: str = Field(description="A brief analysis (2-3 sentences) of whether the figure matches the request and satisfies the criteria")
	passed: bool
	issues: list[str] = Field(description="A list of all the concrete problems found; empty if the figure passes")


SYSTEM_PROMPT = """ You are an expert in Python data visualization. You write runnable Python code using matplotlib.

Rules:
- Use only numpy, pandas, matplotlib. Do not use the internet or other files for writing code.
- Save the plot to OUTPUT_PATH. This variable is already defined. Do not create another file.
- Think through the steps then output the final code in a python block between backticks.
- Use plt.savefig(OUTPUT_PATH, dpi=150, bbox_inches='tight'). Do not call plt.show().
- You may reason step by step first. Your response MUST end with the complete, runnable code in a single ```python block, with nothing after it.
- Only the final code goes inside a ```python block. Any intermediate thoughts or drafts must be in plain text, not in code fences.
- Make a relevant plot: add labels, title, legend and make sure the choice of colours is an understandable one.
"""

GENERATE_PROMPT = """Task: {task}

Write Python code that produces this plot. Reason briefly if needed, then end with the code in a single ```python block. Save it to OUTPUT_PATH."""

REPAIR_PROMPT = """The code you generated failed.

Previous code:
{code}

Error message:
{error}


Understand why the code has failed and generate a correct version for the exact same task.
Reason briefly if needed, then end with the corrected code in a single ```python block."""

REVIEW_PROMPT = """The following figure was requested: {task}

Judge whether the figure is a correct and faithful visualization of this request. Check ONLY the following:
- Chart type: the figure must be the same type of chart that was requested.
- Data: the values shown must match the requested data.
- The image must not be blank or empty.
- The figure must not be misleading (e.g., a truncated axis that hides or exaggerates values).
- Text and labels must be readable and must not overlap or be cut off.

Only report a problem if you can clearly see it in the image. Do not guess or invent issues (for example, do not report spelling mistakes or styling preferences).

Reason briefly, then end your response with a single JSON object in exactly this format:
{{"passed": true, "issues": []}}
Set "passed" to false if any check fails, and list the concrete problems in "issues" (empty list if it passes)."""

ISSUES_PROMPT = """The code ran and produced a figure, but the figure has these problems: {issues}.
Fix the code to address them and satisfy the original task.
Reason briefly, then end with the corrected code in a ```python block."""

def generate(messages: list[dict], temperature) -> str:
	params = SamplingParams(
		temperature = temperature,
		top_p = 0.95,
		max_tokens = 3000,
	)

	outputs = model.chat(messages, params)

	return outputs[0].outputs[0].text

def set_limits():
    resource.setrlimit(resource.RLIMIT_CPU, (5, 5))
    resource.setrlimit(resource.RLIMIT_FSIZE, (50 * 1024**2, 50 * 1024**2))
    resource.setrlimit(resource.RLIMIT_AS, (2 * 1024**3, 2 * 1024**3))


def run_code(code, dir_path=os.path.expanduser('~/qgen/plots')):
	os.makedirs(dir_path, exist_ok=True)

	filename = 'plot-' + str(uuid.uuid4().hex) + '.png'
	# print(file_path)

	plot_path = os.path.join(dir_path, filename)

	header = f"""
	import matplotlib
	matplotlib.use('Agg')

	OUTPUT_PATH = '{plot_path}'
	"""
 
	footer = """
	import matplotlib.pyplot as plt
	_fig = plt.gcf()
	_fig.canvas.draw()
	for _ax in _fig.axes:
		_xbboxes = [lbl.get_window_extent() for lbl in _ax.get_xticklabels()]
		_ybboxes = [lbl.get_window_extent() for lbl in _ax.get_yticklabels()]
		for _i in range(len(_xbboxes)):
			for _j in range(_i + 1, len(_xbboxes)):
				if _xbboxes[_i].overlaps(_xbboxes[_j]):
					print("OVERLAP_X_DETECTED")
		for _i in range(len(_ybboxes)):
			for _j in range(_i + 1, len(_ybboxes)):
				if _ybboxes[_i].overlaps(_ybboxes[_j]):
					print("OVERLAP_Y_DETECTED")
 	"""

	codename = 'gen.py'

	code_path = os.path.join(dir_path, codename)

	with open(code_path, 'w') as f:
		f.write(textwrap.dedent(header) + '\n' + code + '\n' + textwrap.dedent(footer))

	issues = []

	try:
		result = subprocess.run([sys.executable, code_path], capture_output=True, text=True, timeout=5, preexec_fn=set_limits)

		if "OVERLAP_X_DETECTED" in result.stdout:
			issues.append("overlapping labels detected on x axis")
		if "OVERLAP_Y_DETECTED" in result.stdout:
			issues.append("overlapping labels detected on y axis")
  
		if result.returncode:
			return {"image": None, "reason": result.stderr, "success": False}, []
		elif os.path.exists(plot_path):
			return {"image": plot_path, "reason": "no error", "success": True}, issues
		else:
			return {"image": None, "reason": "run but did not produce image", "success": False}, []
	except subprocess.TimeoutExpired as e:
		return {"image": None, "reason": "timeout", "success": False}, []


def extract_code(text):
	matches = re.findall(r"(?<=```python)(.*?)(?=```)", text, flags=re.DOTALL)
	return matches[-1].strip() if matches else text.strip()

def safety_scan(code):
	issues = []

	ALLOWED = {"matplotlib", "numpy", "pandas", "math"}

	try:
		tree = ast.parse(code)
	except:
		return ["code could not be parsed"]

	for node in ast.walk(tree):
		if isinstance(node, ast.Import):
			for alias in node.names:
				main = alias.name.split(".")[0]
				if main not in ALLOWED:
					issues.append(f"forbidden import: {alias.name}")
		elif isinstance(node, ast.ImportFrom):
			module_name = node.module or ""
			main = module_name.split(".")[0]
			if main not in ALLOWED:
				issues.append(f"forbidden import: {module_name}")

	return issues


def check_blank_image(image_path):
	issues = []

	image_array = image.imread(image_path)
	arr_std = image_array.std()

	if arr_std < 0.03:
		issues.append("blank image")
	
	return issues


def check_figure(code):
	issues = []

	if "title(" not in code:
		issues.append("missing title")
	if "xlabel(" not in code:
		issues.append("x-axis is not labeled")
	if "ylabel(" not in code:
		issues.append("y-axis is not labeled")
	return issues


def review_image(image_path, task):
	with open (image_path, "rb") as f:
		encoded_image = base64.b64encode(f.read()).decode()

	messages = [
		{"role": "user", "content": [
		{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{encoded_image}"}},
		{"type": "text", "text": REVIEW_PROMPT.format(task=task)}   
		]}
	]

	structured_outputs = StructuredOutputsParams(json=ReviewOutput.model_json_schema())

	params = SamplingParams(
			temperature = 0.7,
			top_p = 0.95,
			max_tokens = 4000,
			structured_outputs = structured_outputs
	)

	outputs = model.chat(messages, params)

	try:
		return json.loads(outputs[0].outputs[0].text)
	except json.JSONDecodeError:
		return {"passed": True, "issues": []}
		
def generate_plot(task, temperature, max_iter=5):
	messages = [
		{"role": "system", "content": SYSTEM_PROMPT},
		{"role": "user", "content": GENERATE_PROMPT.format(task=task)}
	]

	for attempt in range(max_iter):
		issues = []
		attempt += 1

		reply = generate(messages, temperature)

		code = extract_code(reply)
  
		code_issues = safety_scan(code)
  
		if code_issues:
			messages.extend([
				{"role": "assistant", "content": code},
				{"role": "user", "content": REPAIR_PROMPT.format(code=code, error=", ".join(code_issues))}
			])
			continue

		result, run_issues = run_code(code)

		if result["success"]:
			issues = (
       			check_figure(code) + 
             	review_image(result["image"], task)["issues"] + 
               	check_blank_image(result["image"]) +
				run_issues
			)

			if issues:
				messages.extend([
					{"role": "assistant", "content": code},
					{"role": "user", "content": ISSUES_PROMPT.format(issues=issues)}
				])
			else:
				result.update({"attempts": attempt})
				return result
		else:
			messages.extend([
				{"role": "assistant", "content": code},
				{"role": "user", "content": REPAIR_PROMPT.format(code=code, error=result["reason"])}
			])

	return {"image": None, "reason": "task failed", "success": False, "attempts": attempt}