import os
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load model and tokenizer
model_path = "D1rtyB1rd/Smol-Dirty-Alice"
model = AutoModelForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# Define the stop token and system message
stop_token_id = 2 ## <|im_end|>
system_message = "<|im_start|>system\nYou are a helpful assistant..\n<|im_end|>"

def chat_with_model(prompt_text, stop_token_id, model, tokenizer):
    # Encode the prompt text
    encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False, return_tensors="pt")

    # Generate response
    output_sequences = model.generate(
        input_ids=encoded_prompt,
        max_new_tokens=1024,
        temperature=0.2,
        repetition_penalty=1.2,
        top_k=20,
        top_p=0.9,
        do_sample=True,
        num_return_sequences=1,
        eos_token_id=stop_token_id,
    )

    # Decode the generated sequence
    generated_sequence = output_sequences[0].tolist()
    text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)

    response_text = text[len(prompt_text):].strip()  # Extract only the response text
    return response_text

def build_prompt(conversation_history, user_input):
    """
    Constructs the prompt for the model using conversation history and the latest user input.
    """
    prompt_text = f"{conversation_history}\n<|im_start|>user\n{user_input}\n<|im_end|>\n<|im_start|>assistant\n"
    return prompt_text

def main():
    # Initialize conversation history with the system message
    conversation_history = f"{system_message}\n"

    # Chat loop
    while True:
        user_input = input("User: ")  # Get text input from the user

        # Construct prompt text for model input
        prompt_text = build_prompt(conversation_history, user_input)
        
        response_text = chat_with_model(prompt_text, stop_token_id, model, tokenizer)
        response_text = response_text.replace('<s>', '')

        print(f"\n------\nAlice:\n{response_text}\n------")

        # Update conversation history
        conversation_history += f"<|im_start|>user\n{user_input}\n<|im_end|>\n<|im_start|>assistant\n{response_text}\n<|im_end|>\n"

        # Trim the conversation history to avoid overly long inputs
        if len(conversation_history) > 2048:
            conversation_history = conversation_history[-1024:]

if __name__ == "__main__":
    main()