marcomaccarini/SynthLA

IMPORTS

pip install trl peft torch datasets transformers jupyterlab accelerate tiktoken matplotlib bitsandbytes evaluate scikit-learn

CODE

from huggingface_hub import login
access_token = "secret-token"
login(token=access_token)

import torch
import datasets
from peft import PeftModel
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, AutoModel, TFBertForQuestionAnswering,TFAutoModelWithLMHead

GPU_use = 0
st = "cuda:"+str(GPU_use)
torch.cuda.set_device(GPU_use)

ds = datasets.load_dataset('marcomaccarini/ds_robot_33_large')

trn = ds['train']

base_model = 'meta-llama/Meta-Llama-3-8B'
tokr = AutoTokenizer.from_pretrained(base_model)
model = AutoModelForCausalLM.from_pretrained("marcomaccarini/SynthLA", torch_dtype=torch.bfloat16, device_map=GPU_use,token=access_token)

fmt = """
Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
USER: {}
===
{}
ASSISTANT:"""

def sql_prompt(d): 
    return fmt.format(d["context"], d["question"])
def question(table, quest):
    tst = dict(**trn[8])
    tst['context'] = table
    tst['question'] = quest
    return sql_prompt(tst)

t = 'table([ eof x: 85 y: 179 z: 548, gripper: open , black-cup x: -54 y: -27 z: 80, white-cup x: -5 y: 59 z: 60, box x: -30 y: 34 z: 100, green-cylinder x: 25 y: -3 z: 80 or: 142, green-cube x: -390 y: -490 z: 80 or: 83, grey-cube x: 56 y: -22 z: 80 or: 96, red-cube x: -32 y: 58 z: 80 or: 157, yellow-ball x: -21 y: 30 z: 20 or: 41, banana x: 2 y: 53 z: 20 or: 9, remote x: -48 y: 31 z: 30 or: 69, pen x: -53 y: -59 z: 10 or: 174 ])'
q = 'pick green-cube and place to black-cup'

test = question(t,q)
toks = tokr(test, return_tensors="pt")
res = model.generate(**toks.to(st), max_new_tokens=100, top_p = 0).to('cpu')

print(tokr.batch_decode(res)[0].replace("*","\n"))