jojo-ai-mst commited on
Commit
d9589e2
1 Parent(s): 9ed32b7

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +3 -1
README.md CHANGED
@@ -43,12 +43,14 @@ print(outputs)
43
  ### alternative ways
44
 
45
  ```python
 
 
46
 
47
  model = GPT2LMHeadModel.from_pretrained("jojo-ai-mst/MyanmarGPT")
48
  tokenizer = GPT2Tokenizer.from_pretrained("jojo-ai-mst/MyanmarGPT")
49
 
50
  def generate_text(prompt, max_length=300, temperature=0.8, top_k=50):
51
- input_ids = tokenizer.encode(prompt, return_tensors="pt").cuda()
52
  output = model.generate(
53
  input_ids,
54
  max_length=max_length,
 
43
  ### alternative ways
44
 
45
  ```python
46
+ import torch
47
+ from transformers import AutoTokenizer, AutoModelForCausalLM
48
 
49
  model = GPT2LMHeadModel.from_pretrained("jojo-ai-mst/MyanmarGPT")
50
  tokenizer = GPT2Tokenizer.from_pretrained("jojo-ai-mst/MyanmarGPT")
51
 
52
  def generate_text(prompt, max_length=300, temperature=0.8, top_k=50):
53
+ input_ids = tokenizer.encode(prompt, return_tensors="pt").cuda() # remove .cude() if only cpu
54
  output = model.generate(
55
  input_ids,
56
  max_length=max_length,