from typing import Dict, List, Any
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

class EndpointHandler:
    def __init__(self, path=""):
        # 1. This runs only ONCE when the Endpoint boots up
        print("Loading Ormuri AI into the Cloud GPU...")
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        
        # We load in float16 to save memory and make it run much faster
        self.model = AutoModelForCausalLM.from_pretrained(
            path, 
            device_map="auto", 
            torch_dtype=torch.float16
        )
        print("Model Ready!")

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        # Get the raw message from the website
        user_input = data.pop("inputs", "")
        
        # 1. THE FIX: Format it EXACTLY like your university training data!
        system_prompt = "You are a helpful, accurate, and friendly Ormuri language assistant. Only provide the exact translation. Do not invent words."
        formatted_prompt = f"### Instruction:\n{system_prompt}\n\nUser Question: {user_input}\n\n### Response:\n"
        
        # Convert text to numbers
        input_ids = self.tokenizer(formatted_prompt, return_tensors="pt").input_ids.to(self.model.device)
        
        # 2. THE FIX: Turn down the temperature so it stops hallucinating fake words!
        output_ids = self.model.generate(
            input_ids, 
            max_new_tokens=150, 
            pad_token_id=self.tokenizer.eos_token_id,
            temperature=0.1, # Strictly factual, no guessing
            do_sample=True
        )
        
        # Slice the output to ONLY grab the brand new words after "### Response:\n"
        new_tokens = output_ids[0][input_ids.shape[1]:] 
        
        # Decode back to text
        final_answer = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
        
        # Send ONLY the clean answer back
        return [{"generated_text": final_answer.strip()}]