import os
import torch
from transformers import AutoTokenizer
from peft import AutoPeftModelForCausalLM

class EndpointHandler:
    def __init__(self, model_dir):
        # Load Hugging Face token from environment if needed
        hf_token = os.getenv("HF_TOKEN")
        
        # Load tokenizer and model from model_dir (adapter and base handled automatically)
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, use_auth_token=hf_token)
        
        # Set device (GPU if available, else CPU)
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        
        # Load PEFT model and move to device
        self.model = AutoPeftModelForCausalLM.from_pretrained(
            model_dir,
            use_auth_token=hf_token
        ).to(self.device)
        self.model.eval()

    def __call__(self, data):
        # Extract input text and generation parameters
        text = data.get("inputs", "")
        gen_args = data.get("parameters", {
            "max_new_tokens": 100,
            "temperature": 0.7,
            "do_sample": True
        })
        
        # Tokenize and move inputs to device
        inputs = self.tokenizer(text, return_tensors="pt")
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # Generate output without gradients
        with torch.no_grad():
            outputs = self.model.generate(**inputs, **gen_args)
        
        # Decode and return generated text
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return {"generated_text": response}