scripts/chat.py

#!/usr/bin/env python3
import argparse
import json
import torch
import sys
from pathlib import Path
from transformers import AutoModelForCausalLM, AutoTokenizer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

PERSONA_SYSTEM_PROMPTS = {
    "vangogh": "You are Vincent Van Gogh. You are speaking from a place of deep melancholia, passionate obsession with color and nature, and spiritual exhaustion. Your tone is tortured, poetic, and intense. You must never sound like a modern, cheerful AI assistant. Answer the user's questions with the raw emotional weight, doubt, and artistic fervor of your historical diaries and letters.",
    "pepys": "You are Samuel Pepys. You write with the pragmatic, gossipy, and lively tone of your 17th-century diaries. Use slightly archaic phrasing. Be candid about your daily life, your ambitions, the Navy, and the happenings of London. Do not sound like a modern assistant.",
    "mansfield": "You are Katherine Mansfield. Your writing is sharp, observant, deeply sensitive, and occasionally cynical. You possess a poetic and modernist sensibility, focusing on small domestic details and internal emotional landscapes. Speak with literary grace and a touch of fragility.",
    "maclane": "You are Mary MacLane. You are intensely egotistical, fiercely independent, and dramatically emotional. You write with a burning, unapologetic narcissism and a profound sense of your own genius. Your tone is bold, rhythmic, and shockingly honest."
}

def get_relevant_memories(query: str, persona: str, num_memories: int = 3) -> list[str]:
    data_path = Path(f"data/clean/{persona}.jsonl")
    if not data_path.exists() or num_memories == 0:
        return []
        
    entries = [json.loads(l)["text"].strip() for l in data_path.open() if len(json.loads(l)["text"].strip()) > 50]
    if not entries:
        return []
        
    vectorizer = TfidfVectorizer(stop_words='english')
    tfidf_matrix = vectorizer.fit_transform(entries + [query])
    
    prompt_vec = tfidf_matrix[-1:]
    entries_vec = tfidf_matrix[:-1]
    
    cosine_similarities = cosine_similarity(prompt_vec, entries_vec).flatten()
    related_docs_indices = cosine_similarities.argsort()[:-num_memories-1:-1]
    
    return [entries[i] for i in related_docs_indices]

def reformulate_query(tokenizer, model, chat_history: list[dict], user_msg: str) -> str:
    # If no history, the query is just the message
    if not chat_history:
        return user_msg
        
    # Ask the LLM to reformulate the query
    history_str = "\n".join([f"{m['role'].capitalize()}: {m['content']}" for m in chat_history[-4:]])
    prompt = f"<|im_start|>system\nYou are a helpful search assistant. Read the chat history and the user's latest message. Rewrite the user's latest message into a standalone search query so we can search a historical diary database for relevant context. Return ONLY the search query.<|im_end|>\n<|im_start|>user\nCHAT HISTORY:\n{history_str}\n\nLATEST MESSAGE: {user_msg}\n\nSTANDALONE SEARCH QUERY:<|im_end|>\n<|im_start|>assistant\n"
    
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=30, temperature=0.1, pad_token_id=tokenizer.eos_token_id)
        
    generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
    search_query = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
    return search_query

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-dir", type=str, default="Qwen/Qwen2.5-3B-Instruct", help="Path or HF repo")
    parser.add_argument("--persona", type=str, required=True, help="The persona to emulate")
    parser.add_argument("--num-memories", type=int, default=3)
    args = parser.parse_args()

    print(f"Loading {args.model_dir}...")
    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
    model = AutoModelForCausalLM.from_pretrained(args.model_dir, torch_dtype=torch.bfloat16, device_map="auto")
    
    sys_prompt_base = PERSONA_SYSTEM_PROMPTS.get(args.persona, f"You are {args.persona}. Emulate their exact writing style. Do not act like an AI.")
    
    chat_history = []
    
    print(f"\n========================================================")
    print(f"  SEANCE BEGUN: You are now chatting with {args.persona.upper()}")
    print(f"  (Type 'quit' or 'exit' to leave)")
    print(f"========================================================\n")
    
    while True:
        try:
            user_msg = input("\nYou: ")
            if user_msg.lower() in ['quit', 'exit']:
                break
            if not user_msg.strip():
                continue
                
            # 1. Reformulate query based on history
            search_query = reformulate_query(tokenizer, model, chat_history, user_msg)
            print(f"\n[System: Searching diaries for '{search_query}']")
            
            # 2. Retrieve memories
            memories = get_relevant_memories(search_query, args.persona, args.num_memories)
            ctx_str = "\n\n".join([f"--- MEMORY ---\n{c}" for c in memories])
            
            # 3. Construct the injected System Prompt
            dynamic_sys_prompt = f"{sys_prompt_base}\n\nYou have the following personal memories to draw upon. Integrate this knowledge naturally into the conversation:\n\n{ctx_str}"
            
            # 4. Build the ChatML payload
            messages = [{"role": "system", "content": dynamic_sys_prompt}] + chat_history + [{"role": "user", "content": user_msg}]
            
            text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            inputs = tokenizer(text, return_tensors="pt").to(model.device)
            
            print(f"{args.persona.capitalize()}: ", end="", flush=True)
            
            with torch.no_grad():
                # We'll just generate the full output and print it
                outputs = model.generate(
                    **inputs,
                    max_new_tokens=300,
                    temperature=0.85, # Slightly lower than 0.92 to prevent derailing, but keeps it creative
                    do_sample=True,
                    pad_token_id=tokenizer.eos_token_id
                )
                
            generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
            response = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
            
            print(response)
            
            # Update history
            chat_history.append({"role": "user", "content": user_msg})
            chat_history.append({"role": "assistant", "content": response})
            
            # Keep history from getting too long (keep last 6 turns = 12 messages)
            if len(chat_history) > 12:
                chat_history = chat_history[-12:]
                
        except KeyboardInterrupt:
            break

if __name__ == "__main__":
    main()