vanta-research/apollo-astralis-4b

1"""2Apollo-Astralis V1 4B - Example Usage3 4This script demonstrates how to use Apollo-Astralis V1 4B with Transformers.5"""6 7from transformers import AutoModelForCausalLM, AutoTokenizer8import torch9 10def load_model(model_name="VANTA-Research/apollo-astralis-v1-4b"):11    """Load Apollo-Astralis model and tokenizer."""12    print(f"Loading {model_name}...")13    14    tokenizer = AutoTokenizer.from_pretrained(15        model_name,16        trust_remote_code=True17    )18    19    model = AutoModelForCausalLM.from_pretrained(20        model_name,21        torch_dtype=torch.bfloat16,22        device_map="auto",23        trust_remote_code=True24    )25    26    print("Model loaded successfully!")27    return model, tokenizer28 29def generate_response(model, tokenizer, user_message, system_prompt=None):30    """Generate a response from Apollo."""31    if system_prompt is None:32        system_prompt = "You are Apollo-Astralis V1, a warm and enthusiastic reasoning assistant."33    34    messages = [35        {"role": "system", "content": system_prompt},36        {"role": "user", "content": user_message}37    ]38    39    # Apply chat template40    text = tokenizer.apply_chat_template(41        messages,42        tokenize=False,43        add_generation_prompt=True44    )45    46    # Tokenize47    inputs = tokenizer([text], return_tensors="pt").to(model.device)48    49    # Generate50    outputs = model.generate(51        **inputs,52        max_new_tokens=512,53        temperature=0.7,54        top_p=0.9,55        do_sample=True,56        repetition_penalty=1.0557    )58    59    # Decode60    response = tokenizer.decode(61        outputs[0][inputs['input_ids'].shape[1]:],62        skip_special_tokens=True63    )64    65    return response66 67def main():68    # Load model69    model, tokenizer = load_model()70    71    # Example 1: Celebration72    print("\n" + "="*60)73    print("Example 1: Celebration Response")74    print("="*60)75    user_msg = "I just got my first job as a software engineer!"76    print(f"\nUser: {user_msg}")77    response = generate_response(model, tokenizer, user_msg)78    print(f"\nApollo: {response}")79    80    # Example 2: Problem-solving81    print("\n" + "="*60)82    print("Example 2: Problem-Solving")83    print("="*60)84    user_msg = "What's the best way to learn machine learning?"85    print(f"\nUser: {user_msg}")86    response = generate_response(model, tokenizer, user_msg)87    print(f"\nApollo: {response}")88    89    # Example 3: Mathematical reasoning90    print("\n" + "="*60)91    print("Example 3: Mathematical Reasoning")92    print("="*60)93    user_msg = "If a train travels 120 km in 1.5 hours, what's its average speed?"94    print(f"\nUser: {user_msg}")95    response = generate_response(model, tokenizer, user_msg)96    print(f"\nApollo: {response}")97 98if __name__ == "__main__":99    main()100