# Step 1 (Run once)
!pip install -U transformers

# Step 2
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

# Using pipeline
pipe = pipeline("text-generation", model="sargurun16/VCoder")

messages = [
    {"role": "user", "content": "Who are you?"},
]

print(pipe(messages))

# Step 3 & 4
model_name = "sargurun16/VCoder"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

prompt = "write a python code to merge 3 arrays"

inputs = tokenizer(
    prompt,
    return_tensors="pt"
)

outputs = model.generate(
    **inputs,
    max_new_tokens=200
)

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)