import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def interactive_chat():
    # Relative OS path for open-source robustness
    base_dir = os.path.dirname(os.path.abspath(__file__))
    model_dir = os.path.join(base_dir, "sail_5b_hf_model")
    
    print(f"Loading tokenizer from {model_dir}...")
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
    except Exception as e:
        print(f"Failed to load tokenizer: {e}")
        return

    print(f"Loading 350M SAIL model from {model_dir}...")
    try:
        # Load model. Since it's untrained, the output will be gibberish, but it will confirm the architecture works!
        model = AutoModelForCausalLM.from_pretrained(
            model_dir,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto",
            trust_remote_code=True
        )
    except Exception as e:
        print(f"Failed to load model: {e}")
        return
    
    print("\n========================================================")
    print("  SAIL 350M Foundational Network - Interactive Chat")
    print("  NOTE: This model is currently UNTRAINED blank weights.")
    print("  Outputs will be completely random until pre-training.")
    print("========================================================")
    
    while True:
        try:
            user_input = input("\nYou: ")
            if user_input.lower() in ['quit', 'exit', 'stop']:
                break
            if not user_input.strip():
                continue
            
            # Format message
            messages = [{"role": "user", "content": user_input}]
            # Apply chat template
            prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            
            inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
            
            # Generate response
            with torch.no_grad():
                outputs = model.generate(
                    **inputs,
                    max_new_tokens=150,
                    do_sample=True,
                    temperature=0.7,
                    top_p=0.9,
                    pad_token_id=tokenizer.pad_token_id,
                    eos_token_id=tokenizer.eos_token_id,
                )
            
            # Decode only the generated part
            input_length = inputs["input_ids"].shape[1]
            response_tokens = outputs[0][input_length:]
            response = tokenizer.decode(response_tokens, skip_special_tokens=True)
            
            print(f"\nSAIL AI: {response}")
            
        except KeyboardInterrupt:
            break
        except Exception as e:
            print(f"\nError during generation: {e}")

if __name__ == "__main__":
    interactive_chat()