diff options
| author | sigoden <sigoden@gmail.com> | 2024-06-12 19:17:40 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-06-12 19:17:40 +0800 |
| commit | 64982b4510e38153885bfd0a78c250110b3e03c5 (patch) | |
| tree | 896eb39e51b17f7a4a52901a21c4f1cbb78abb90 /src/rag/splitter/language.rs | |
| parent | 492b006db7e5e40b2672e28dc93cc1bc68a7cf38 (diff) | |
| download | aichat-64982b4510e38153885bfd0a78c250110b3e03c5.tar.gz | |
feat: rag splitter supports languages (#593)
Diffstat (limited to 'src/rag/splitter/language.rs')
| -rw-r--r-- | src/rag/splitter/language.rs | 235 |
1 files changed, 235 insertions, 0 deletions
diff --git a/src/rag/splitter/language.rs b/src/rag/splitter/language.rs new file mode 100644 index 0000000..20722cf --- /dev/null +++ b/src/rag/splitter/language.rs @@ -0,0 +1,235 @@ +#[derive(PartialEq, Eq, Hash)] +pub enum Language { + Cpp, + Go, + Java, + Js, + Php, + Proto, + Python, + Rst, + Ruby, + Rust, + Scala, + Swift, + Markdown, + Latex, + Html, + Sol, +} + +impl Language { + pub fn separators(&self) -> Vec<&str> { + match self { + Language::Cpp => vec![ + "\nclass ", + "\nvoid ", + "\nint ", + "\nfloat ", + "\ndouble ", + "\nif ", + "\nfor ", + "\nwhile ", + "\nswitch ", + "\ncase ", + "\n\n", + "\n", + " ", + "", + ], + Language::Go => vec![ + "\nfunc ", + "\nvar ", + "\nconst ", + "\ntype ", + "\nif ", + "\nfor ", + "\nswitch ", + "\ncase ", + "\n\n", + "\n", + " ", + "", + ], + Language::Java => vec![ + "\nclass ", + "\npublic ", + "\nprotected ", + "\nprivate ", + "\nstatic ", + "\nif ", + "\nfor ", + "\nwhile ", + "\nswitch ", + "\ncase ", + "\n\n", + "\n", + " ", + "", + ], + Language::Js => vec![ + "\nfunction ", + "\nconst ", + "\nlet ", + "\nvar ", + "\nclass ", + "\nif ", + "\nfor ", + "\nwhile ", + "\nswitch ", + "\ncase ", + "\ndefault ", + "\n\n", + "\n", + " ", + "", + ], + Language::Php => vec![ + "\nfunction ", + "\nclass ", + "\nif ", + "\nforeach ", + "\nwhile ", + "\ndo ", + "\nswitch ", + "\ncase ", + "\n\n", + "\n", + " ", + "", + ], + Language::Proto => vec![ + "\nmessage ", + "\nservice ", + "\nenum ", + "\noption ", + "\nimport ", + "\nsyntax ", + "\n\n", + "\n", + " ", + "", + ], + Language::Python => vec!["\nclass ", "\ndef ", "\n\tdef ", "\n\n", "\n", " ", ""], + Language::Rst => vec![ + "\n===\n", "\n---\n", "\n***\n", "\n.. ", "\n\n", "\n", " ", "", + ], + Language::Ruby => vec![ + "\ndef ", + "\nclass ", + "\nif ", + "\nunless ", + "\nwhile ", + "\nfor ", + "\ndo ", + "\nbegin ", + "\nrescue ", + "\n\n", + "\n", + " ", + "", + ], + Language::Rust => vec![ + "\nfn ", "\nconst ", "\nlet ", "\nif ", "\nwhile ", "\nfor ", "\nloop ", + "\nmatch ", "\nconst ", "\n\n", "\n", " ", "", + ], + Language::Scala => vec![ + "\nclass ", + "\nobject ", + "\ndef ", + "\nval ", + "\nvar ", + "\nif ", + "\nfor ", + "\nwhile ", + "\nmatch ", + "\ncase ", + "\n\n", + "\n", + " ", + "", + ], + Language::Swift => vec![ + "\nfunc ", + "\nclass ", + "\nstruct ", + "\nenum ", + "\nif ", + "\nfor ", + "\nwhile ", + "\ndo ", + "\nswitch ", + "\ncase ", + "\n\n", + "\n", + " ", + "", + ], + Language::Markdown => vec![ + "\n## ", + "\n### ", + "\n#### ", + "\n##### ", + "\n###### ", + "```\n\n", + "\n\n***\n\n", + "\n\n---\n\n", + "\n\n___\n\n", + "\n\n", + "\n", + " ", + "", + ], + Language::Latex => vec![ + "\n\\chapter{", + "\n\\section{", + "\n\\subsection{", + "\n\\subsubsection{", + "\n\\begin{enumerate}", + "\n\\begin{itemize}", + "\n\\begin{description}", + "\n\\begin{list}", + "\n\\begin{quote}", + "\n\\begin{quotation}", + "\n\\begin{verse}", + "\n\\begin{verbatim}", + "\n\\begin{align}", + "$$", + "$", + "\n\n", + "\n", + " ", + "", + ], + Language::Html => vec![ + "<body>", "<div>", "<p>", "<br>", "<li>", "<h1>", "<h2>", "<h3>", "<h4>", "<h5>", + "<h6>", "<span>", "<table>", "<tr>", "<td>", "<th>", "<ul>", "<ol>", "<header>", + "<footer>", "<nav>", "<head>", "<style>", "<script>", "<meta>", "<title>", " ", "", + ], + Language::Sol => vec![ + "\npragma ", + "\nusing ", + "\ncontract ", + "\ninterface ", + "\nlibrary ", + "\nconstructor ", + "\ntype ", + "\nfunction ", + "\nevent ", + "\nmodifier ", + "\nerror ", + "\nstruct ", + "\nenum ", + "\nif ", + "\nfor ", + "\nwhile ", + "\ndo while ", + "\nassembly ", + "\n\n", + "\n", + " ", + "", + ], + } + } +} |