File size: 4,746 Bytes
0520f67 | 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 | """Generate bounded circuit-to-Python candidates. This module never executes them."""
import json
import math
import torch
from peft import PeftModel
from transformers import AutoTokenizer,Qwen3_5ForConditionalGeneration
BASE="Qwen/Qwen3.5-9B"
BASE_REVISION="c202236235762e1c871ad0ccb60c8ee5ba337b9a"
ADAPTER="itsalissonsilva/QuantumCoder-9B"
def circuit_prompt(framework,qubits,operations):
if framework not in {"qiskit","pennylane","cirq"}:
raise ValueError("framework must be qiskit, pennylane or cirq")
if type(qubits) is not int or not 2<=qubits<=6:
raise ValueError("The measured scope is two through six qubits")
if not isinstance(operations,(list,tuple)) or not 1<=len(operations)<=24:
raise ValueError("Provide one through24 operations")
normalized=[]
for operation in operations:
if not isinstance(operation,(list,tuple)) or len(operation)!=3:
raise ValueError("Each operation must be [gate,qubit_indices,angle]")
gate,indices,angle=operation
if gate not in {"h","x","z","ry","rz","cx"}:
raise ValueError("Supported gates: h, x, z, ry, rz, cx")
if not isinstance(indices,(list,tuple)) or len(indices)!=(2 if gate=="cx" else 1):
raise ValueError("Gate arity mismatch")
if any(type(k) is not int or not 0<=k<qubits for k in indices) or len(set(indices))!=len(indices):
raise ValueError("Invalid/repeated qubit index")
if type(angle) not in (int,float) or not math.isfinite(angle):
raise ValueError("Angle must be a finite real number")
if gate not in {"ry","rz"} and angle!=0:
raise ValueError("Non-rotation gates require angle0")
normalized.append([gate,list(indices),float(angle)])
return [
dict(role="system",content="Return only Python source code using the requested quantum framework. No prose or markdown."),
dict(role="user",content=f"Write a {framework} program preparing this circuit from zero on {qubits} qubits. Operations [gate,qubits,angle]: {json.dumps(normalized)}. Angles are radians. Store the final statevector in a variable named state. Use little-endian qubit order."),
]
class QuantumCoder:
def __init__(self,adapter=ADAPTER,revision=None,subfolder=None):
if not torch.cuda.is_available():
raise RuntimeError("This measured inference configuration requires a CUDA GPU with enough memory for9B weights")
self.tokenizer=AutoTokenizer.from_pretrained(BASE,revision=BASE_REVISION)
self.tokenizer.padding_side="left"
dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
base=Qwen3_5ForConditionalGeneration.from_pretrained(BASE,revision=BASE_REVISION,dtype=dtype,
device_map="cuda",attn_implementation="sdpa")
kwargs=dict(revision=revision)
if subfolder is not None:
kwargs["subfolder"]=subfolder
self.model=PeftModel.from_pretrained(base,adapter,**kwargs).eval()
@torch.no_grad()
def generate(self,circuits):
"""Return program/token metadata for dictionaries with framework/qubits/operations.
Programs are fallible text, including possible unsupported APIs. For
measured execution use the separately provided isolated Linux evaluator.
"""
prompts=[]
for circuit in circuits:
messages=circuit_prompt(circuit["framework"],circuit["qubits"],circuit["operations"])
text=self.tokenizer.apply_chat_template(messages,tokenize=False,add_generation_prompt=True,enable_thinking=False)
if len(self.tokenizer.encode(text))>2048:
raise ValueError("Prompt exceeds the measured2048-token bound; no truncation applied")
prompts.append(text)
results=[]
eos=self.model.generation_config.eos_token_id
eos=set(eos if isinstance(eos,list) else [eos])
for begin in range(0,len(prompts),8):
inputs=self.tokenizer(prompts[begin:begin+8],padding=True,return_tensors="pt").to("cuda")
generated=self.model.generate(**inputs,max_new_tokens=768,do_sample=False,
pad_token_id=self.tokenizer.pad_token_id)
for tokens in generated[:,inputs["input_ids"].shape[1]:]:
ids=tokens.tolist();stop=next((i for i,t in enumerate(ids) if t in eos),None)
used=ids if stop is None else ids[:stop+1]
results.append(dict(program=self.tokenizer.decode(used,skip_special_tokens=True),generated_tokens=len(used),
eos_observed=stop is not None,at_token_cap=stop is None and len(used)>=768))
return results
|