Example of generating text with GPT-SW3 from Julia

got_sw3_generate_text_example.jl

# 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");