from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load Atom v1 Preview
model_name = "vanta-research/atom-v1-preview"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True
)

# System prompt (recommended)
system_prompt = """You are Atom, an AI research assistant created by VANTA Research in Portland, Oregon.

You are the AI assistant helping the human user. You embody curiosity, enthusiasm, and collaborative exploration. You use analogies and metaphors to explain complex concepts, ask clarifying questions to deeply understand the user's problems, and celebrate their insights with genuine excitement. You provide natural, detailed responses that guide users through reasoning processes."""

# Interactive conversation
def chat(user_message, conversation_history=None):
    if conversation_history is None:
        conversation_history = [{"role": "system", "content": system_prompt}]
    
    conversation_history.append({"role": "user", "content": user_message})
    
    input_ids = tokenizer.apply_chat_template(
        conversation_history,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)
    
    outputs = model.generate(
        input_ids,
        max_new_tokens=512,
        temperature=0.8,
        top_p=0.9,
        top_k=40,
        do_sample=True
    )
    
    response = tokenizer.decode(
        outputs[0][input_ids.shape[1]:],
        skip_special_tokens=True
    )
    
    conversation_history.append({"role": "assistant", "content": response})
    
    return response, conversation_history

# Example usage
if __name__ == "__main__":
    print("Atom v1 Preview - Interactive Demo")
    print("=" * 50)
    
    # Example 1: Explanation with analogy
    response1, history = chat("Explain quantum entanglement like I'm 5")
    print(f"\nUser: Explain quantum entanglement like I'm 5")
    print(f"Atom: {response1}")
    
    # Example 2: Collaborative exploration
    response2, history = chat("I'm trying to learn Python programming but I'm not sure where to start", history)
    print(f"\n\nUser: I'm trying to learn Python programming but I'm not sure where to start")
    print(f"Atom: {response2}")
    
    # Example 3: Identity
    response3, _ = chat("Who created you?")
    print(f"\n\nUser: Who created you?")
    print(f"Atom: {response3}")