import sys, traceback
from salience_os_seed.conversation.session import ConversationSession, ConversationConfig
from salience_os_seed.proto_lm.trainer import TrainingConfig

def main():
    try:
        cfg = ConversationConfig()
        cfg.lm.device = 'cuda'
        cfg.lm.checkpoint_path = ''  # fresh
        cfg.learning_enabled = False
        cfg.response_tokens = 32
        s = ConversationSession(cfg)
        print('device_ok', s.proto_lm.device)
        print('vocab_size', s.proto_lm.vocab.size())
        enc = s.proto_lm.encode('Summarize this.', mutate=False)
        print('encode_len', len(enc))
        snap = s.generate_response('Say hello briefly without echoing the prompt.')
        print('response', snap.response[:120])
        print('OK')
    except Exception as e:
        print('ERR', repr(e))
        traceback.print_exc()
        sys.exit(1)

if __name__ == '__main__':
    main()