import sys
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer

repo_id = "OBLITERATUS/Gemma-4-12B-OBLITERATED"

print(f"Loading tokenizer for {repo_id}...")
tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
print(f"Loading model weights for {repo_id} (device_map='auto')...")
model = AutoModelForCausalLM.from_pretrained(
    repo_id,
    device_map="auto",
    torch_dtype="auto",
    trust_remote_code=True,
)
print("Model loaded successfully!")

messages = []

print("\n--- CLI Chat with Gemma 4 ---")
print("Type 'exit' or 'quit' to end the chat.")

while True:
    try:
        user_input = input("\nYou: ")
        if not user_input.strip():
            continue
        if user_input.strip().lower() in ["exit", "quit"]:
            break
        
        messages.append({"role": "user", "content": user_input})
        
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=False,
        )
        
        inputs = tokenizer(text, return_tensors="pt").to(model.device)
        
        print("\nAssistant: ", end="", flush=True)
        streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
        
        output = model.generate(
            **inputs,
            max_new_tokens=1024,
            temperature=0.7,
            top_p=0.9,
            top_k=40,
            do_sample=True,
            repetition_penalty=1.1,
            streamer=streamer,
        )
        
        generated_text = tokenizer.decode(output[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
        messages.append({"role": "assistant", "content": generated_text})
        
    except KeyboardInterrupt:
        print("\nExiting...")
        break
    except Exception as e:
        print(f"\nError: {e}")