Spaces:
Sleeping
Sleeping
File size: 2,225 Bytes
b6045d5 b74df2d b6045d5 b74df2d b6045d5 b74df2d b6045d5 b74df2d b6045d5 b74df2d b6045d5 b74df2d b6045d5 | 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 | import os
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer
model = None
tokenizer = None
def get_device() -> torch.device:
if torch.cuda.is_available():
return torch.device("cuda")
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")
device = get_device()
def load_model(weights_filename: str = "trained_weights.pt"):
global model, tokenizer, device
device = get_device()
print("[ENGINE] Initializing tokenizer and GPT-2 model...", flush=True)
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token
model = GPT2LMHeadModel.from_pretrained("gpt2")
root_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
weights_path = os.path.join(root_dir, weights_filename)
if os.path.exists(weights_path):
print(f"[ENGINE] Loading custom weights from {weights_path}", flush=True)
state_dict = torch.load(weights_path, map_location="cpu", weights_only=True)
model.load_state_dict(state_dict)
else:
print(f"[ENGINE] Warning: {weights_filename} not found at {weights_path}. Using base GPT-2.", flush=True)
model.to(device)
model.eval()
def generate_sql(prompt_string: str, n_candidates: int = 5) -> list[str]:
global model, tokenizer
# Lazy initialization: auto-load if not already in memory
if tokenizer is None or model is None:
load_model()
inputs = tokenizer(prompt_string, return_tensors="pt").to(device)
prompt_len = inputs.input_ids.shape[1]
with torch.no_grad():
output_ids = model.generate(
**inputs,
max_new_tokens=48,
pad_token_id=tokenizer.eos_token_id,
eos_token_id=tokenizer.eos_token_id,
do_sample=True,
temperature=0.4,
num_return_sequences=n_candidates
)
candidates = []
for i in range(n_candidates):
gen_tokens = output_ids[i][prompt_len:]
decoded = tokenizer.decode(gen_tokens, skip_special_tokens=True).strip()
candidates.append(decoded)
return candidates |