|
|
|
|
|
""" |
|
|
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() |
|
|
|
|
|
|
|
|
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() |
|
|
|