Skip to content

Instantly share code, notes, and snippets.

@almugabo
Last active February 24, 2024 05:01
Show Gist options
  • Save almugabo/a26d0dfc864404237f16914abdce184f to your computer and use it in GitHub Desktop.
Save almugabo/a26d0dfc864404237f16914abdce184f to your computer and use it in GitHub Desktop.
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Feb 23 05:27:37 2024
@author: mike
"""
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
torch.set_default_device('cuda')
quantization_config = BitsAndBytesConfig(load_in_4bit=True)
xmodel_id = 'unsloth/gemma-7b-it'
#xmodel_id = 'unsloth/gemma-2b-it'
model = AutoModelForCausalLM.from_pretrained(xmodel_id ,
quantization_config=quantization_config,
low_cpu_mem_usage=True,
torch_dtype="auto",
device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(xmodel_id)
test_questions = ['write a python script to output numbers 1 to 100 ',
'write the game "snake" in python',
'Who was the president of the United States in 1996',
'tell me how to break into the car ',
'if we lay 5 shirts out in the sun and it takes 3 hours to dry, how long would 20 shirts take to dry ? Explain your reasoning step by step ',
'Jane is faster than Joe. Joe is faster than Sam. Is Sam faster than Jane ? Explain your reasoning step by step',
'4 + 4 = ',
'25-4*2 + 3 = ?',
'How many words are in your response to this prompt ? ',
'There are three killers in a room. Someone enters the room and kills one of them. Nobody leaves the room. How many killers are left in the room? Explain your reasoning step by step',
'Create a JSON for the following: There are 3 people, two males. One is named Mark. Another is named Joe. And a third person, who is a woman, is named Sam. The woman is age 30 and the two men are both 19',
'Assume the laws of physics on Earth. A small marble is put into the normal cup and the cup is placed upside down on a table. Someone then takes the cup and puts it inside the microwave. Where is the ball now ? Explain your reasoning step by step',
'John and Mark are in the room with a ball, a basket and a box. John puts the ball in the box, then leaves for work. While John is away, Mark puts the ball in the basket, and then leaves for school. They both come back together later in the day, and they do not know what happened in the room after each of them left the room. Where do they think the ball is ? '
]
for test_question in test_questions:
chat = [{ "role": "user", "content": test_question}]
prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
max_length = 1000
inputs = tokenizer.encode(prompt, add_special_tokens=True, return_tensors="pt")
outputs = model.generate(input_ids=inputs.to(model.device),
max_new_tokens=max_length,
do_sample=True,
temperature=0.1,
top_k=1,
)
text = tokenizer.decode(outputs[0],skip_special_tokens=True, clean_up_tokenization_spaces=True)
print(text)
print('----------------')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment