132 lines
4.4 KiB
Python
132 lines
4.4 KiB
Python
import json
|
||
from abc import ABC
|
||
from typing import List
|
||
from langchain.docstore.document import Document
|
||
from langchain.document_loaders.base import BaseLoader
|
||
|
||
|
||
class Person:
|
||
def __init__(self, name, age):
|
||
self.name = name
|
||
self.age = age
|
||
|
||
|
||
class Dialogue:
|
||
"""
|
||
Build an abstract dialogue model using classes and methods to represent different dialogue elements.
|
||
This class serves as a fundamental framework for constructing dialogue models.
|
||
"""
|
||
|
||
def __init__(self, file_path: str):
|
||
self.file_path = file_path
|
||
self.turns = []
|
||
|
||
def add_turn(self, turn):
|
||
"""
|
||
Create an instance of a conversation participant
|
||
:param turn:
|
||
:return:
|
||
"""
|
||
self.turns.append(turn)
|
||
|
||
def parse_dialogue(self):
|
||
"""
|
||
The parse_dialogue function reads the specified dialogue file and parses each dialogue turn line by line.
|
||
For each turn, the function extracts the name of the speaker and the message content from the text,
|
||
creating a Turn instance. If the speaker is not already present in the participants dictionary,
|
||
a new Person instance is created. Finally, the parsed Turn instance is added to the Dialogue object.
|
||
|
||
Please note that this sample code assumes that each line in the file follows a specific format:
|
||
<speaker>:\r\n<message>\r\n\r\n. If your file has a different format or includes other metadata,
|
||
you may need to adjust the parsing logic accordingly.
|
||
"""
|
||
participants = {}
|
||
speaker_name = None
|
||
message = None
|
||
|
||
with open(self.file_path, encoding='utf-8') as file:
|
||
lines = file.readlines()
|
||
for i, line in enumerate(lines):
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
|
||
if speaker_name is None:
|
||
speaker_name, _ = line.split(':', 1)
|
||
elif message is None:
|
||
message = line
|
||
if speaker_name not in participants:
|
||
participants[speaker_name] = Person(speaker_name, None)
|
||
|
||
speaker = participants[speaker_name]
|
||
turn = Turn(speaker, message)
|
||
self.add_turn(turn)
|
||
|
||
# Reset speaker_name and message for the next turn
|
||
speaker_name = None
|
||
message = None
|
||
|
||
def display(self):
|
||
for turn in self.turns:
|
||
print(f"{turn.speaker.name}: {turn.message}")
|
||
|
||
def export_to_file(self, file_path):
|
||
with open(file_path, 'w', encoding='utf-8') as file:
|
||
for turn in self.turns:
|
||
file.write(f"{turn.speaker.name}: {turn.message}\n")
|
||
|
||
def to_dict(self):
|
||
dialogue_dict = {"turns": []}
|
||
for turn in self.turns:
|
||
turn_dict = {
|
||
"speaker": turn.speaker.name,
|
||
"message": turn.message
|
||
}
|
||
dialogue_dict["turns"].append(turn_dict)
|
||
return dialogue_dict
|
||
|
||
def to_json(self):
|
||
dialogue_dict = self.to_dict()
|
||
return json.dumps(dialogue_dict, ensure_ascii=False, indent=2)
|
||
|
||
def participants_to_export(self):
|
||
"""
|
||
participants_to_export
|
||
:return:
|
||
"""
|
||
participants = set()
|
||
for turn in self.turns:
|
||
participants.add(turn.speaker.name)
|
||
return ', '.join(participants)
|
||
|
||
|
||
class Turn:
|
||
def __init__(self, speaker, message):
|
||
self.speaker = speaker
|
||
self.message = message
|
||
|
||
|
||
class DialogueLoader(BaseLoader, ABC):
|
||
"""Load dialogue."""
|
||
|
||
def __init__(self, file_path: str):
|
||
"""Initialize with dialogue."""
|
||
self.file_path = file_path
|
||
dialogue = Dialogue(file_path=file_path)
|
||
dialogue.parse_dialogue()
|
||
self.dialogue = dialogue
|
||
|
||
def load(self) -> List[Document]:
|
||
"""Load from dialogue."""
|
||
documents = []
|
||
participants = self.dialogue.participants_to_export()
|
||
|
||
for turn in self.dialogue.turns:
|
||
metadata = {"source": f"Dialogue File：{self.dialogue.file_path},"
|
||
f"speaker：{turn.speaker.name}，"
|
||
f"participant：{participants}"}
|
||
turn_document = Document(page_content=turn.message, metadata=metadata.copy())
|
||
documents.append(turn_document)
|
||
|
||
return documents
|