60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
import threading
|
|
from stream import Stream
|
|
|
|
VALID_ROLES = ['user', 'assistant', 'system']
|
|
|
|
class LlamaMessage:
|
|
def __init__(self, role, content):
|
|
if role not in VALID_ROLES:
|
|
raise Exception("message role must be one of: " + ",".join(VALID_ROLES))
|
|
self.role = role
|
|
self.content = content
|
|
|
|
|
|
class Models:
|
|
llama_8b = "meta/meta-llama-3-8b-instruct"
|
|
llama_70b = "meta/meta-llama-3-70b-instruct"
|
|
|
|
async def replicate_llama_completion(app, messages, model=None, temperature=0.75, max_tokens=500):
|
|
if model is None:
|
|
model = Models.llama_8b
|
|
|
|
# Llama instruct
|
|
full_prompt = "<|begin_of_text|>"
|
|
for message in messages:
|
|
full_prompt += f"<|start_header_id|>{message.role}<|end_header_id|>\n\n{message.content}<|eot_id|>"
|
|
full_prompt += f"<|start_header_id|>assistant<|end_header_id|>"
|
|
|
|
content_stream = Stream(app)
|
|
thread = threading.Thread(
|
|
target=run_replicate_inference_loop,
|
|
kwargs={
|
|
'replicate_client': app.replicate_client,
|
|
'model': model,
|
|
'prompt': full_prompt,
|
|
'content_stream': content_stream,
|
|
'temperature': temperature,
|
|
'max_tokens': max_tokens,
|
|
}
|
|
)
|
|
thread.start()
|
|
response = await content_stream.content_wait()
|
|
return response.lstrip()
|
|
|
|
|
|
def run_replicate_inference_loop(replicate_client, model, prompt, content_stream, temperature, max_tokens):
|
|
print("START INFERENCE")
|
|
completion = replicate_client.run(
|
|
model,
|
|
input={
|
|
"prompt": prompt,
|
|
"temperature": temperature,
|
|
"max_new_tokens": max_tokens,
|
|
"system_prompt": "" # We handle system prompt in the formatted prompt
|
|
}
|
|
)
|
|
for event in completion:
|
|
if event:
|
|
content_stream.feed(str(event))
|
|
print("END INFERENCE")
|
|
content_stream.close() |