Views
No views yet
1import argparse, re, torch
2from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, BitsAndBytesConfig
3from peft import PeftModel
4
5STOP_STR = "</OpenScenario>"
6
7class StopOnSubstrings(StoppingCriteria):
8 def __init__(self, tok, substrings, start_len=0):
9 self.tok = tok
10 self.substrings = substrings
11 self.start_len = start_len
12
13 def set_start_len(self, n):
14 self.start_len = n
15
16 def __call__(self, input_ids, scores, **kw):
17 gen_ids = input_ids[0, self.start_len:].tolist()
18 if not gen_ids:
19 return False
20 text = self.tok.decode(gen_ids, skip_special_tokens=True)
21 return any(s in text for s in self.substrings)
22
23def load_model_and_tokenizer(base_path: str, lora_id: str, use_4bit: bool):
24 # Tokenizer: prefer the LoRA repo (to inherit chat template), fallback to base
25 try:
26 tok = AutoTokenizer.from_pretrained(lora_id, use_fast=True)
27 except Exception:
28 tok = AutoTokenizer.from_pretrained(base_path, use_fast=True)
29 tok.padding_side = "right"
30 if tok.pad_token is None:
31 tok.pad_token = tok.eos_token
32
33 quant_cfg = None
34 device_map = None
35 torch_dtype = torch.bfloat16
36
37 if use_4bit:
38 quant_cfg = BitsAndBytesConfig(
39 load_in_4bit=True,
40 bnb_4bit_quant_type="nf4",
41 bnb_4bit_use_double_quant=True,
42 bnb_4bit_compute_dtype=torch.bfloat16,
43 )
44 device_map = "auto"
45
46 model = AutoModelForCausalLM.from_pretrained(
47 base_path,
48 torch_dtype=None if use_4bit else torch_dtype,
49 quantization_config=quant_cfg,
50 device_map=device_map,
51 attn_implementation="sdpa",
52 )
53 model = PeftModel.from_pretrained(model, lora_id)
54 model.config.use_cache = True
55 model.eval()
56 model.config.pad_token_id = tok.pad_token_id
57 model.config.eos_token_id = tok.eos_token_id
58 return model, tok
59
60def generate_xosc(model, tok, template_path: str, system: str, user: str, max_new: int = 4000) -> str:
61 tmpl = open(template_path, "r", encoding="utf-8").read()
62 prompt = tmpl.format(system=system.strip(), user=user.strip())
63 if prompt.endswith("[/INST]"):
64 prompt += "\n"
65 enc = tok(prompt, return_tensors="pt")
66 enc = {k: v.to(model.device) for k, v in enc.items()}
67 stopper = StopOnSubstrings(tok, [STOP_STR])
68 stopper.set_start_len(enc["input_ids"].shape[1])
69
70 with torch.inference_mode():
71 out = model.generate(
72 **enc,
73 max_new_tokens=max_new,
74 do_sample=False,
75 temperature=0.0,
76 top_p=1.0,
77 pad_token_id=tok.pad_token_id,
78 eos_token_id=tok.eos_token_id,
79 stopping_criteria=StoppingCriteriaList([stopper]),
80 return_dict_in_generate=True,
81 )
82 gen_ids = out.sequences[0, enc["input_ids"].shape[1]:]
83 txt = tok.decode(gen_ids, skip_special_tokens=True)
84 m = re.search(r"<OpenScenario\b.*?</OpenScenario>", txt, flags=re.DOTALL)
85 return m.group(0).strip() if m else txt
86
87if __name__ == "__main__":
88 ap = argparse.ArgumentParser()
89 ap.add_argument("--base")
90 ap.add_argument("--lora", default="anto0699/Prompt2OpenSCENARIO-CodeLlama13B-LoRA")
91 ap.add_argument("--template", default=r"prompt_templates\codellama_inst.txt")
92 ap.add_argument("--use_4bit", action="store_true", help="Enable 4-bit loading (Option B). If omitted, use full-precision (Option A).")
93 ap.add_argument("--max_new_tokens", type=int, default=4000)
94 ap.add_argument("--system", default=(
95 "Act as an OpenSCENARIO 1.0 generator for ADS testing in CARLA. "
96 "I will give you a scene description in English and you must return one valid .xosc file, XML only, "
97 "encoded in UTF-8, starting with <OpenScenario> and ending with </OpenScenario>. The file must be schema-compliant, "
98 "and executable in CARLA without modifications. The scenario must include: the map (<RoadNetwork>), "
99 "<Environment> with <TimeOfDay> and <Weather>, exactly one ego vehicle, any other entities with unique names, "
100 "initial positions using <WorldPosition>, and a valid <Storyboard> with deterministic triggers/events/actions. "
101 "Use realistic defaults if details are missing (no randomness); no comments or extra text."
102 ))
103 ap.add_argument("--user", default="Write a minimal scenario with only one ego vehicle in Towns04, sunny environment.")
104 args = ap.parse_args()
105
106 model, tok = load_model_and_tokenizer(args.base, args.lora, args.use_4bit)
107 xosc = generate_xosc(model, tok, args.template, args.system, args.user, args.max_new_tokens)
108 print(xosc)1<OpenScenario>
2 ...
3 <Entities>
4 <ScenarioObject name="ego_vehicle">
5 <Vehicle name="vehicle.lincoln.mkz2017" vehicleCategory="car">
6 ...
7 </Vehicle>
8 </ScenarioObject>
9 </Entities>
10 <Storyboard>
11 ...
12 </Storyboard>
13</OpenScenario>q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj)1@misc{prompt2openscenario2025,
2 title = {Empty Title},
3 author = {No Authors},
4 year = {2025},
5 publisher = {Hugging Face},
6 howpublished = {\url{https://huggingface.co/anto0699/Prompt2OpenSCENARIO-CodeLlama13B-LoRA}}
7}