| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| from utils.config import Config |
| from utils.logger import Logger |
|
|
| logger = Logger.get_logger(__name__) |
|
|
|
|
| class ModelGenerator: |
| """ |
| Singleton class responsible for generating text using a specified language model. |
| |
| This class initializes a language model and tokenizer, and provides methods |
| to generate text and extract code blocks from generated text. |
| |
| Attributes: |
| device (torch.device): Device to run the model on (CPU or GPU). |
| model (AutoModelForCausalLM): Language model for text generation. |
| tokenizer (AutoTokenizer): Tokenizer corresponding to the language model. |
| |
| Methods: |
| acceptTextGenerator(self, visitor, *args, **kwargs): |
| Accepts a visitor to generates text based on the input provided with the model generator. |
| acceptExtractCodeBlock(self, visitor, *args, **kwargs): |
| Accepts a visitor to extract code blocks from the output text. |
| """ |
| _instance = None |
| _format_data_time = "%Y-%m-%d %H:%M:%S" |
|
|
| def __new__(cls, model_name=Config.read('app', 'model')): |
| if cls._instance is None: |
| cls._instance = super(ModelGenerator, cls).__new__(cls) |
| cls._instance._initialize(model_name) |
| return cls._instance |
|
|
| def _initialize(self, model_name): |
| self.device = torch.device( |
| "cuda" if torch.cuda.is_available() else "cpu") |
| self.model = AutoModelForCausalLM.from_pretrained( |
| model_name).to(self.device) |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
|
| def acceptTextGenerator(self, visitor, *args, **kwargs): |
| return visitor.visit(self, *args, **kwargs) |
|
|
| def acceptExtractCodeBlock(self, visitor, *args, **kwargs): |
| return visitor.visit(self, *args, **kwargs) |
|
|
