Fill-Mask
Transformers
code
iamthe66epitaph commited on
Commit
9c8351b
·
verified ·
1 Parent(s): cb20617

Upload baby_ai.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. baby_ai.py +19 -0
baby_ai.py ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoModelForCausalLM, AutoTokenizer
2
+ # Load a tiny model
3
+ model_name = "distilgpt2"
4
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
5
+ model = AutoModelForCausalLM.from_pretrained(model_name)
6
+
7
+ print("Hello, I'm your baby AI! Type 'quit' to exit.")
8
+
9
+ while True:
10
+ user_input = input("You: ")
11
+ if user_input.lower() == "quit":
12
+ break
13
+
14
+ inputs = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
15
+ outputs = model.generate(inputs, max_length=100, pad_token_id=tokenizer.eos_token_id)
16
+
17
+ response = tokenizer.decode(outputs[0], skip_special_tokens=True)
18
+ print("AI:", response)
19
+