wassname/Open-Assistant
1
0
Fork 0
mirror of https://github.com/wassname/Open-Assistant.git synced 2026-08-13 11:50:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

Message tree state machine (#555)

Browse Source 
* add query_incomplete_rankings()

* Add SQL queries for TreeManager task selection

* first working version of TreeManager.next_task()

* remove old generate_task(), add mandatory_labels to text_labels task

* Add ConversationMessage list to Ranking tasks

* add more sophisticated sql queries to find extendible trees

* add TreeManager.query_extendible_parents()

* fix task validation, seed data insertion (reviewed)

* provide user for task selection in text-frontend

* enter 'growing' state

* enter 'aborted_low_grade' state

* enter 'ranking' state

* check tree 'growing' state upon relpy insertion

* exclude user from labeling their own messages (added DEBUG_ALLOW_SELF_LABELING setting)

* add DEBUG_ALLOW_SELF_LABELING to docker-compose.yaml

* fix ranking submission

* add query_tree_ranking_results()

* add ranked_message_ids to RankingReactionPayload

* fix reply_messages instead of prompt_messages

* incorment 'ranking_count' of ranked replies

* added logic to check_condition_for_scoring_state

* changes to msg_tree_state_machine

* pre-commit changes

* enter 'ready_for_scoring' state

* re-add HF embedding call (lost during merge)

* use prepare_conversation() helper for seed-data creation

* Partially add user specified task selection

Co-authored-by: Daniel Hug <danielpatrickhug@gmail.com>
This commit is contained in:
Andreas Köpf
2023-01-11 10:54:03 +01:00
committed by GitHub
co-authored by Daniel Hug
parent 23ff01c603
commit 14fa08e2e7
19 changed files with 1212 additions and 323 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/oasst_backend/api/v1/tasks.py
+8 -216
View File
@@ -1,15 +1,12 @@
import random
from typing import Any, Optional, Tuple
from typing import Any
from uuid import UUID
from fastapi import APIRouter, Depends
from fastapi.security.api_key import APIKey
from loguru import logger
from oasst_backend.api import deps
from oasst_backend.api.v1.utils import prepare_conversation
from oasst_backend.config import settings
from oasst_backend.prompt_repository import PromptRepository, TaskRepository
from oasst_backend.utils.hugging_face import HfEmbeddingModel, HfUrl, HuggingFaceAPI
from oasst_backend.tree_manager import TreeManager, TreeManagerConfiguration
from oasst_shared.exceptions import OasstError, OasstErrorCode
from oasst_shared.schemas import protocol as protocol_schema
from sqlmodel import Session
@@ -18,160 +15,6 @@ from starlette.status import HTTP_204_NO_CONTENT
router = APIRouter()
def generate_task(
request: protocol_schema.TaskRequest, pr: PromptRepository
) -> Tuple[protocol_schema.Task, Optional[UUID], Optional[UUID]]:
message_tree_id = None
parent_message_id = None
match request.type:
case protocol_schema.TaskRequestType.random:
logger.info("Frontend requested a random task.")
disabled_tasks = (
protocol_schema.TaskRequestType.random,
protocol_schema.TaskRequestType.summarize_story,
protocol_schema.TaskRequestType.rate_summary,
)
candidate_tasks = set(protocol_schema.TaskRequestType).difference(disabled_tasks)
request.type = random.choice(tuple(candidate_tasks)).value
return generate_task(request, pr)
# AKo: Summary tasks are currently disabled/supported, we focus on the conversation tasks.
# case protocol_schema.TaskRequestType.summarize_story:
# logger.info("Generating a SummarizeStoryTask.")
# task = protocol_schema.SummarizeStoryTask(
# story="This is a story. A very long story. So long, it needs to be summarized.",
# )
# case protocol_schema.TaskRequestType.rate_summary:
# logger.info("Generating a RateSummaryTask.")
# task = protocol_schema.RateSummaryTask(
# full_text="This is a story. A very long story. So long, it needs to be summarized.",
# summary="This is a summary.",
# scale=protocol_schema.RatingScale(min=1, max=5),
# )
case protocol_schema.TaskRequestType.initial_prompt:
logger.info("Generating an InitialPromptTask.")
task = protocol_schema.InitialPromptTask(
hint="Ask the assistant about a current event." # this is optional
)
case protocol_schema.TaskRequestType.prompter_reply:
logger.info("Generating a PrompterReplyTask.")
messages = pr.fetch_random_conversation("assistant")
task_messages = [
protocol_schema.ConversationMessage(
text=msg.text,
is_assistant=(msg.role == "assistant"),
message_id=msg.id,
front_end_id=msg.frontend_message_id,
)
for msg in messages
]
task = protocol_schema.PrompterReplyTask(conversation=protocol_schema.Conversation(messages=task_messages))
message_tree_id = messages[-1].message_tree_id
parent_message_id = messages[-1].id
case protocol_schema.TaskRequestType.assistant_reply:
logger.info("Generating a AssistantReplyTask.")
messages = pr.fetch_random_conversation("prompter")
task_messages = [
protocol_schema.ConversationMessage(
text=msg.text,
is_assistant=(msg.role == "assistant"),
message_id=msg.id,
front_end_id=msg.frontend_message_id,
)
for msg in messages
]
task = protocol_schema.AssistantReplyTask(conversation=protocol_schema.Conversation(messages=task_messages))
message_tree_id = messages[-1].message_tree_id
parent_message_id = messages[-1].id
case protocol_schema.TaskRequestType.rank_initial_prompts:
logger.info("Generating a RankInitialPromptsTask.")
messages = pr.fetch_random_initial_prompts()
task = protocol_schema.RankInitialPromptsTask(prompts=[msg.text for msg in messages])
case protocol_schema.TaskRequestType.rank_prompter_replies:
logger.info("Generating a RankPrompterRepliesTask.")
conversation, replies = pr.fetch_multiple_random_replies(message_role="assistant")
task_messages = [
protocol_schema.ConversationMessage(
text=p.text,
is_assistant=(p.role == "assistant"),
message_id=p.id,
front_end_id=p.frontend_message_id,
)
for p in conversation
]
replies = [p.text for p in replies]
task = protocol_schema.RankPrompterRepliesTask(
conversation=protocol_schema.Conversation(
messages=task_messages,
),
replies=replies,
)
case protocol_schema.TaskRequestType.rank_assistant_replies:
logger.info("Generating a RankAssistantRepliesTask.")
conversation, replies = pr.fetch_multiple_random_replies(message_role="prompter")
task_messages = [
protocol_schema.ConversationMessage(
text=p.text,
is_assistant=(p.role == "assistant"),
message_id=p.id,
front_end_id=p.frontend_message_id,
)
for p in conversation
]
replies = [p.text for p in replies]
task = protocol_schema.RankAssistantRepliesTask(
conversation=prepare_conversation(conversation),
replies=replies,
)
case protocol_schema.TaskRequestType.label_initial_prompt:
logger.info("Generating a LabelInitialPromptTask.")
message = pr.fetch_random_initial_prompts(1)[0]
task = protocol_schema.LabelInitialPromptTask(
message_id=message.id,
prompt=message.text,
valid_labels=list(map(lambda x: x.value, protocol_schema.TextLabel)),
)
case protocol_schema.TaskRequestType.label_prompter_reply:
logger.info("Generating a LabelPrompterReplyTask.")
conversation, messages = pr.fetch_multiple_random_replies(max_size=1, message_role="assistant")
message = messages[0]
task = protocol_schema.LabelPrompterReplyTask(
message_id=message.id,
conversation=prepare_conversation(conversation),
reply=message.text,
valid_labels=list(map(lambda x: x.value, protocol_schema.TextLabel)),
)
case protocol_schema.TaskRequestType.label_assistant_reply:
logger.info("Generating a LabelAssistantReplyTask.")
conversation, messages = pr.fetch_multiple_random_replies(max_size=1, message_role="prompter")
message = messages[0]
task = protocol_schema.LabelAssistantReplyTask(
message_id=message.id,
conversation=prepare_conversation(conversation),
reply=message.text,
valid_labels=list(map(lambda x: x.value, protocol_schema.TextLabel)),
)
case _:
raise OasstError("Invalid request type", OasstErrorCode.TASK_INVALID_REQUEST_TYPE)
logger.info(f"Generated {task=}.")
return task, message_tree_id, parent_message_id
@router.post(
"/",
response_model=protocol_schema.AnyTask,
@@ -193,7 +36,9 @@ def request_task(
try:
pr = PromptRepository(db, api_client, client_user=request.user)
task, message_tree_id, parent_message_id = generate_task(request, pr)
tree_manager_config = TreeManagerConfiguration()
tm = TreeManager(db, pr, tree_manager_config)
task, message_tree_id, parent_message_id = tm.next_task(request.type)
pr.task_repository.store_task(task, message_tree_id, parent_message_id, request.collective)
except OasstError:
@@ -268,63 +113,10 @@ async def tasks_interaction(
try:
pr = PromptRepository(db, api_client, client_user=interaction.user)
tree_manager_config = TreeManagerConfiguration()
tm = TreeManager(db, pr, tree_manager_config)
return await tm.handle_interaction(interaction)
match type(interaction):
case protocol_schema.TextReplyToMessage:
logger.info(
f"Frontend reports text reply to {interaction.message_id=} with {interaction.text=} by {interaction.user=}."
)
# here we store the text reply in the database
newMessage = pr.store_text_reply(
text=interaction.text,
frontend_message_id=interaction.message_id,
user_frontend_message_id=interaction.user_message_id,
)
if not settings.DEBUG_SKIP_EMBEDDING_COMPUTATION:
try:
hugging_face_api = HuggingFaceAPI(
f"{HfUrl.HUGGINGFACE_FEATURE_EXTRACTION.value}/{HfEmbeddingModel.MINILM.value}"
)
embedding = await hugging_face_api.post(interaction.text)
pr.insert_message_embedding(
message_id=newMessage.id, model=HfEmbeddingModel.MINILM.value, embedding=embedding
)
except OasstError:
logger.error(
f"Could not fetch embbeddings for text reply to {interaction.message_id=} with {interaction.text=} by {interaction.user=}."
)
return protocol_schema.TaskDone()
case protocol_schema.MessageRating:
logger.info(
f"Frontend reports rating of {interaction.message_id=} with {interaction.rating=} by {interaction.user=}."
)
# here we store the rating in the database
pr.store_rating(interaction)
return protocol_schema.TaskDone()
case protocol_schema.MessageRanking:
logger.info(
f"Frontend reports ranking of {interaction.message_id=} with {interaction.ranking=} by {interaction.user=}."
)
# TODO: check if the ranking is valid
pr.store_ranking(interaction)
# here we would store the ranking in the database
return protocol_schema.TaskDone()
case protocol_schema.TextLabels:
logger.info(
f"Frontend reports labels of {interaction.message_id=} with {interaction.labels=} by {interaction.user=}."
)
# Labels are implicitly validated when converting str -> TextLabel
# So no need for explicit validation here
pr.store_text_labels(interaction)
return protocol_schema.TaskDone()
case _:
raise OasstError("Invalid response type.", OasstErrorCode.TASK_INVALID_RESPONSE_TYPE)
except OasstError:
raise
except Exception:
backend/oasst_backend/api/v1/utils.py
+12 -11
View File
@@ -18,19 +18,20 @@ def prepare_message_list(messages: list[Message]) -> list[protocol.Message]:
return [prepare_message(m) for m in messages]
def prepare_conversation(messages: list[Message]) -> protocol.Conversation:
conv_messages = []
for message in messages:
conv_messages.append(
protocol.ConversationMessage(
text=message.text,
is_assistant=(message.role == "assistant"),
message_id=message.id,
frontend_message_id=message.frontend_message_id,
)
def prepare_conversation_message_list(messages: list[Message]) -> list[protocol.ConversationMessage]:
return [
protocol.ConversationMessage(
text=message.text,
is_assistant=(message.role == "assistant"),
message_id=message.id,
frontend_message_id=message.frontend_message_id,
)
for message in messages
]
return protocol.Conversation(messages=conv_messages)
def prepare_conversation(messages: list[Message]) -> protocol.Conversation:
return protocol.Conversation(messages=prepare_conversation_message_list(messages))
def prepare_tree(tree: list[Message], tree_id: UUID) -> protocol.MessageTree: