обновлен формат генерации датасета
This commit is contained in:
+1
-68
@@ -1,6 +1,5 @@
|
||||
import os
|
||||
import webbrowser
|
||||
import pandas as pd
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic_ai.models.openai import OpenAIChatModel
|
||||
@@ -13,7 +12,6 @@ from pydantic_ai.providers.openai_codex import (
|
||||
|
||||
from .config import CodexConfig, LLMConfig, OpenAiConfig
|
||||
from .credentials import FileCodexCredentionalSource
|
||||
from .schemas import BaseDatasetRecord, ChatRecord, ChatMessage
|
||||
|
||||
|
||||
class LLMModelProvider:
|
||||
@@ -22,7 +20,7 @@ class LLMModelProvider:
|
||||
def create_openai_compatible(model_name: str,
|
||||
api_key: str,
|
||||
base_url: str
|
||||
) -> OpenAICodexModel:
|
||||
) -> OpenAIChatModel:
|
||||
|
||||
provider = OpenAIProvider(base_url=base_url,
|
||||
api_key=api_key)
|
||||
@@ -84,68 +82,3 @@ class LLMModelProvider:
|
||||
)
|
||||
|
||||
raise ValueError(f"Unsupported LLM config: {type(config)}")
|
||||
|
||||
class DatasetWriter:
|
||||
def __init__(self,
|
||||
dataset_record: BaseDatasetRecord,
|
||||
base_path: str | Path):
|
||||
|
||||
|
||||
self._dataset_record = dataset_record
|
||||
self._path = Path(base_path)
|
||||
self._path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def __check_file_exists(self, filename :str)-> bool:
|
||||
file_path = self._path / filename
|
||||
return file_path.exists()
|
||||
|
||||
def to_csv(self,
|
||||
filename: str) -> None:
|
||||
|
||||
exists=self.__check_file_exists(filename)
|
||||
|
||||
df = pd.json_normalize(
|
||||
self._dataset_record.model_dump()
|
||||
)
|
||||
|
||||
df.to_csv(self._path / filename,
|
||||
mode="a" if exists else "w",
|
||||
header=not exists,
|
||||
index=False,
|
||||
encoding="utf-8")
|
||||
|
||||
def to_chat_record(self) -> ChatRecord:
|
||||
response = self._dataset_record.response
|
||||
prompts = self._dataset_record.prompts
|
||||
|
||||
assistant_content = (
|
||||
response.assistant_content
|
||||
or response.model_response
|
||||
)
|
||||
|
||||
return ChatRecord(
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="system",
|
||||
content=prompts.system_prompt,
|
||||
),
|
||||
ChatMessage(
|
||||
role="user",
|
||||
content=prompts.user_prompt,
|
||||
),
|
||||
ChatMessage(
|
||||
role="assistant",
|
||||
content=assistant_content,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
def to_jsonl(self,
|
||||
filename: str) -> None:
|
||||
|
||||
chat_record = self.to_chat_record()
|
||||
path = self._path/filename
|
||||
|
||||
with path.open("a", encoding='utf-8') as file:
|
||||
file.write(chat_record.model_dump_json())
|
||||
file.write("\n")
|
||||
Reference in New Issue
Block a user