|
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) |
|
|
|
|
|
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 |
|
|
