#!/usr/bin/env python3
"""
Example usage of the SmolLM2 Help Bot
"""

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

def main():
    print("Loading SmolLM2 Help Bot...")
    
    tokenizer = AutoTokenizer.from_pretrained("./")
    model = AutoModelForCausalLM.from_pretrained(
        "./",
        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        device_map="auto" if torch.cuda.is_available() else None
    )
    
    def get_help_advice(question):
        prompt = f"Human: {question}\n\nAssistant:"
        inputs = tokenizer(prompt, return_tensors="pt")
        
        if torch.cuda.is_available():
            inputs = {k: v.cuda() for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = model.generate(
                inputs["input_ids"],
                attention_mask=inputs["attention_mask"],
                max_new_tokens=150,
                do_sample=True,
                temperature=0.8,
                top_p=0.9,
                pad_token_id=tokenizer.eos_token_id
            )
        
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response[len(prompt):].strip()
    
    # Interactive mode
    print("\nSmolLM2 Help Bot is ready! Ask any question and get motivational advice.")
    print("Type 'quit' to exit.\n")
    
    while True:
        question = input("You: ")
        if question.lower() == 'quit':
            break
        
        advice = get_help_advice(question)
        print(f"\nHelp Bot: {advice}\n")
        print("-" * 80)

if __name__ == "__main__":
    main()