"""
===========================================================
 LLM Loader (Transformers / Hugging Face)
===========================================================
"""

# ---------------- IMPORTS ----------------
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from .base_model import BaseModel

# ---------------- CLASS ----------------
class DarkNeuronLLM(BaseModel):
    """Handles fine-tuned transformer models."""

    def __init__(self, model_name: str, device=None):
        self.model_name = model_name
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.tokenizer = None
        self.model = None

    # ---------------- LOAD MODEL ----------------
    def load_model(self):
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModelForCausalLM.from_pretrained(self.model_name).to(self.device)
        return self

    # ---------------- GENERATE ----------------
    def generate(self, text, max_length=200):
        inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
        outputs = self.model.generate(**inputs, max_length=max_length)
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)

    def predict(self, text):
        # Required by BaseModel abstract class
        return self.generate(text)