Spaces:
Sleeping
Sleeping
| 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 |