phase2/dl_generator.py

import sys


class DLCodeGenerator:
    _instance = None
    _initialized = False

    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self, model_name: str = "Salesforce/codet5-base", enable_dl: bool = True):
        if self._initialized:
            return
        self.model_name = model_name
        self.enable_dl = enable_dl
        self.tokenizer = None
        self.model = None
        self.is_mock = not enable_dl
        self._initialized = True

    def _ensure_loaded(self) -> bool:
        """Lazily load the CodeT5 model and tokenizer only when generate is first called."""
        if not self.enable_dl:
            return False
        if self.model is not None and self.tokenizer is not None:
            return True
        print(f"Lazy loading DL model: {self.model_name}...")
        try:
            from transformers import RobertaTokenizer, AutoModelForSeq2SeqLM, AutoTokenizer

            try:
                self.tokenizer = RobertaTokenizer.from_pretrained(
                    self.model_name, extra_ids=0, additional_special_tokens=[], local_files_only=False
                )
            except Exception:
                self.tokenizer = AutoTokenizer.from_pretrained(
                    self.model_name, local_files_only=False
                )

            self.model = AutoModelForSeq2SeqLM.from_pretrained(
                self.model_name, local_files_only=False
            )
            self.is_mock = False
            print("DL Model loaded successfully.")
            return True
        except Exception as e:
            print(
                f"Note: DL Model '{self.model_name}' skipped ({e}). Using ultra-fast AST rule generator."
            )
            self.is_mock = True
            return False

    def generate(
        self,
        code: str,
        prompt: str,
        temperature: float = 0.7,
        max_length: int = 512,
    ) -> str:
        if not self._ensure_loaded():
            return code

        try:
            input_text = f"{prompt}:\n{code}"
            input_ids = self.tokenizer(
                input_text,
                return_tensors="pt",
                truncation=True,
                max_length=max_length,
            ).input_ids
            outputs = self.model.generate(
                input_ids,
                max_new_tokens=max_length,
                temperature=temperature,
                do_sample=True if temperature > 0 else False,
                top_p=0.95,
                num_return_sequences=1,
            )
            return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        except Exception as e:
            print(f"DL Generation error: {e}. Falling back to source code.")
            return code