mirror of https://github.com/sigoden/aichat
feat: rag splitter supports languages (#593)
parent
492b006db7
commit
64982b4510
@ -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",
|
||||||
|
" ",
|
||||||
|
"",
|
||||||
|
],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
Loading…
Reference in New Issue