starchat implementation (solved)

[jloganolson]

UPDATE: I can't figure out how to delete this so I guess it's just an example of using starchat? If folks have improvements, please do weigh in...

import torch
from transformers import AutoTokenizer,AutoModelForCausalLM
import guidance
model_path = "HuggingFaceH4/starchat-alpha"

class StarcoderChat(guidance.llms.Transformers):
    def __init__(self, model_path, **kwargs):
        tokenizer = AutoTokenizer.from_pretrained(model_path, device_map='auto')
        model = AutoModelForCausalLM.from_pretrained(model_path, device_map='auto', torch_dtype=torch.bfloat16)
        super().__init__(model, tokenizer=tokenizer, device_map='auto', **kwargs)
    
    @staticmethod
    def role_start(role):
        return "<|"+role+"|>"
    
    @staticmethod
    def role_end(role):
        return '<|end|>' 

guidance.llm = StarcoderChat(model_path)

prompt = guidance('''
{{#system~}}
You are a helpful and terse assistant.
{{~/system}}
{{#user~}}
How do you print something in python?
{{~/user}}
{{#assistant~}}
{{gen 'answer'}}
{{~/assistant}}''')
prompt()