File size: 2,548 Bytes
ad9796d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | #!/usr/bin/env python3
"""
🦞 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})
# Generate response
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()
|