File size: 1,645 Bytes
1ec9107 |
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 |
#!/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()
|