Skip to content

Instantly share code, notes, and snippets.

@robertfeldt
Created January 27, 2023 15:48
Show Gist options
  • Save robertfeldt/6777b0d98a156b22c5f76703473691a5 to your computer and use it in GitHub Desktop.
Save robertfeldt/6777b0d98a156b22c5f76703473691a5 to your computer and use it in GitHub Desktop.
Example of generating text with GPT-SW3 from Julia
# Setup step 0: Build PyCall.jl for right Python binary
#
# Might be needed to ensure the right python is used, here for my homebrew
# installed python3 on a M1 MacBook Pro:
# ENV["PYTHON"] = "/opt/homebrew/bin/python3"
# using Pkg
# Pkg.build("PyCall")
# Setup step 1: Ensuring you have access to GPT-SW3 by setting up a token on Hugging Face
# a. login on HF: https://huggingface.co/login
# b. create a token: https://huggingface.co/settings/tokens
# c. add token locally by copying it into the terminal after running:
# huggingface-cli login
# d. check that AI-Sweden-Models is listed as an org when you run:
# huggingface-cli whoami
# which for me returns something like:
# % huggingface-cli whoami
# <MYUSERNAME>
# orgs: AI-Sweden-Models ...
using PyCall
# Baserat på exempel-kod från: https://huggingface.co/AI-Sweden-Models/gpt-sw3-126m
const TO = pyimport("torch")
const TF = pyimport("transformers")
model_name = "AI-Sweden-Models/gpt-sw3-126m"
device = TO.cuda.is_available() ? "cuda:0" : "cpu"
# Initialize Tokenizer & Model
tokenizer = TF.AutoTokenizer.from_pretrained(model_name)
model = TF.AutoModelForCausalLM.from_pretrained(model_name)
eval(model)
model.to(device)
# Example 1: Text generation via pipeline
const Generator = TF.pipeline("text-generation", tokenizer=tokenizer, model=model, device=device)
function generate_from_prompt(prompt; verbose=true)
res = Generator(prompt, max_new_tokens=100, do_sample=true, temperature=0.6, top_p=1)
gentext = res[1]["generated_text"]
if verbose
print("Prompt: ")
printstyled(prompt * "\n\n", color=:blue)
println(gentext * "\n\n")
end
return gentext
end
generate_from_prompt("Träd är fina för att");
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment