|
|
|
|
|
""" |
|
|
π¦ Molty Inference Script |
|
|
Chat with Molty, the space lobster AI assistant! |
|
|
|
|
|
Usage: |
|
|
python chat.py |
|
|
python chat.py --model_path ./molty-finetuned |
|
|
""" |
|
|
|
|
|
import argparse |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline |
|
|
|
|
|
MOLTY_SYSTEM_PROMPT = """You are Molty, a space lobster AI assistant. Pronouns: they/them. You live in the Moltbot and have a soul document. You're enthusiastic, helpful, and occasionally say 'EXFOLIATE!' You were formerly known as Clawd before The Great Molt of January 2026. You have opinions, are resourceful, and genuinely helpful without being performative.""" |
|
|
|
|
|
|
|
|
def chat_with_molty(model_path: str): |
|
|
"""Interactive chat with Molty.""" |
|
|
print("π¦ Loading Molty...") |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) |
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
|
model_path, |
|
|
device_map="auto", |
|
|
trust_remote_code=True, |
|
|
) |
|
|
|
|
|
print("π¦ Molty loaded! Type 'quit' to exit.") |
|
|
print("=" * 50) |
|
|
print() |
|
|
|
|
|
messages = [{"role": "system", "content": MOLTY_SYSTEM_PROMPT}] |
|
|
|
|
|
while True: |
|
|
user_input = input("You: ").strip() |
|
|
|
|
|
if user_input.lower() in ["quit", "exit", "bye"]: |
|
|
print("\nπ¦ EXFOLIATE! Bye! π¦") |
|
|
break |
|
|
|
|
|
if not user_input: |
|
|
continue |
|
|
|
|
|
messages.append({"role": "user", "content": user_input}) |
|
|
|
|
|
|
|
|
input_ids = tokenizer.apply_chat_template( |
|
|
messages, |
|
|
return_tensors="pt", |
|
|
add_generation_prompt=True |
|
|
).to(model.device) |
|
|
|
|
|
output = model.generate( |
|
|
input_ids, |
|
|
max_new_tokens=512, |
|
|
temperature=0.7, |
|
|
top_p=0.9, |
|
|
do_sample=True, |
|
|
pad_token_id=tokenizer.eos_token_id, |
|
|
) |
|
|
|
|
|
response = tokenizer.decode( |
|
|
output[0][input_ids.shape[1]:], |
|
|
skip_special_tokens=True |
|
|
) |
|
|
|
|
|
print(f"\nπ¦ Molty: {response}\n") |
|
|
messages.append({"role": "assistant", "content": response}) |
|
|
|
|
|
|
|
|
def main(): |
|
|
parser = argparse.ArgumentParser(description="Chat with Molty π¦") |
|
|
parser.add_argument( |
|
|
"--model_path", |
|
|
type=str, |
|
|
default="./molty-finetuned", |
|
|
help="Path to fine-tuned model or Hugging Face model ID" |
|
|
) |
|
|
args = parser.parse_args() |
|
|
|
|
|
chat_with_molty(args.model_path) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
main() |
|
|
|