vicgalle

/

gpt2-open-instruct-v1

like 31

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("vicgalle/gpt2-open-instruct-v1")
model = AutoModelForCausalLM.from_pretrained("vicgalle/gpt2-open-instruct-v1")

# Move model to GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Define the system prompt
system_prompt = """Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
Pretend you are an alien visiting Earth. Write three opinions you believe.

### Response:
"""

# Tokenize the input
inputs = tokenizer(system_prompt, return_tensors="pt").to(device)

# Generate a response
output = model.generate(
    inputs["input_ids"],
    max_length=150,            # Maximum length of the generated sequence
    num_return_sequences=1,    # Number of sequences to return
    temperature=0.7,           # Controls the randomness of the output
    top_k=50,                  # Limits the sampling pool to the top k tokens
    top_p=0.9,                 # Uses nucleus sampling to consider the top p probability mass
    pad_token_id=tokenizer.eos_token_id  # Ensures padding token is set
)

# Decode the output
response = tokenizer.decode(output[0], skip_special_tokens=True)
print(response)