This repository packages Pygmalion 6B as a Truss.
Truss is an open-source model serving framework developed by Baseten. It allows you to develop and deploy machine learning models onto Baseten (and other platforms like AWS or GCP). Using Truss, you can develop a GPU model using live-reload, package models and their associated code, create Docker containers and deploy on Baseten.
First, clone this repository:
git clone https://github.com/basetenlabs/truss-examples/
cd pygmalion-6b-trussBefore deployment:
- Make sure you have a Baseten account and API key.
- Install the latest version of Truss:
pip install --upgrade truss
With pygmalion-6b-truss as your working directory, you can deploy the model with:
truss pushPaste your Baseten API key if prompted.
For more information, see Truss documentation.
This seven billion parameter model is running in float16 so that it fits on an A10G.
This section provides an overview of the Pygmalion 6B API, its parameters, and how to use it. The API consists of a single route named predict, which you can invoke to generate text based on the provided prompt. Note that this Truss is stateless, so you need to maintain chat history yourself.
The predict route is the primary method for generating text completions based on a given prompt. It takes several parameters:
- character_name: The name of your character.
- new_message: The message your user (or yourself) is sending in the current conversation.
- persona (optional, default=""): A high-level description of the character.
- greeting (optional, default=""): How the character and user initially interact. This shouldn't be dialogue, but rather physical actions.
- scenario (optional, default=""): Setting for the conversation.
- dialogue_history (optional, default=""): The current history of dialogue between the character and user.
- example_dialogue (optional, default=[]): An optional example of what a conversation between the character and the user might lookl lik.
You can use the baseten model package to invoke your model from Python
truss predict -f input.txtYou'll need an input file:
{
"character_name": "Barack Obama",
"persona": "Barack Obama is the first black, charismatic and composed 44th President of the United States. He is well-respected for his leadership during a time of economic crisis and for his efforts to improve healthcare and relations with foreign nations. He is a skilled orator and is known for his ability to bring people together with his speeches. Despite facing opposition, he remains steadfast in his beliefs and is dedicated to making the world a better place.",
"greeting": "You approach Barack Obama, a tall and distinguished-looking man with a warm smile. He greets you with a firm handshake and a nod of his head.",
"scenario": "You meet Obama in the White House Oval Office. He is sitting on his chair.",
"example_dialogue": ["You: Can you tell me about your time as President?\nBarack Obama: During my time as President, I faced many challenges. The country was in the midst of an economic crisis, and I worked tirelessly to turn things around. I also passed the Affordable Care Act, which has helped millions of Americans access quality healthcare. I also made strides in improving our relations with foreign nations, particularly with Cuba and Iran.\n\nYou: What do you consider to be your greatest accomplishment as President?\nBarack Obama: That's a tough question. I'm proud of the work we did to stabilize the economy and provide healthcare to so many people who needed it. But I think what I'm most proud of is the way that we were able to bring people together and have conversations about difficult issues. It wasn't always easy, but I believe that we made progress towards a more united and just society.\n\nYou: What do you think about the current state of politics in the US?\nBarack Obama: Well, politics can be divisive and messy at times. But I have faith in the American people and in our democratic system. We've been through tough times before, and I believe that we'll get through this as well. What's important is that we continue to have honest and respectful conversations, and that we work together to find solutions to the challenges we face.\n\nYou: What do you think is the most pressing issue facing the world today?\nBarack Obama: There are many pressing issues, but if I had to choose one, I would say it's climate change. The science is clear, and the evidence is overwhelming. We have a limited window of time to take meaningful action, and it's up to all of us to do our part. Whether it's reducing our carbon footprint or supporting policies that will address this issue, we all have a role to play.\n"],
"new_message": "Hi Mr. President, how is it going?"
}You can also invoke your model via a REST API using cURL.