restore github issue loader, refactoring

README.md

@@ -11,3 +11,38 @@ license: mit
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# import GitHub issues
+
+## export from github
+first, generate token on: https://github.com/settings/tokens
+
+```
+$ git clone https://github.com/kazamori/github-api-tools 
+$ pip install -e ./github-api-tools 
+$ export GITHUB_API_TOKEN="********"
+$ gh-cli-issues --repository <org/repo>
+$ ls <repo>-issues.json
+```
+
+## import from json
+
+```
+$ python store.py -l github_issue <index> ../<repo>-issues.json 
+```
+
+# import Wiki Pages
+
+## export from somewhere
+
+create `pages.json` like:
+```json
+{"id": <page_id>, "title": <page title>, "content": <page body>, "ctime": ..., "user": <name>, "url": "https:..."}
+{"title": ...}
+```
+
+## import from json
+
+```
+$ python store.py -l wikipage <index> ../pages.json 
+```

app.py

@@ -1,6 +1,7 @@
 from time import time
 from datetime import datetime
 from typing import Iterable
+
 import streamlit as st
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
@@ -12,8 +13,9 @@ from qdrant_client.http.models import Filter, FieldCondition, MatchValue, Range
 from langchain.chains import RetrievalQA
 from openai.error import InvalidRequestError
 from langchain.chat_models import ChatOpenAI
+
 from config import DB_CONFIG
-from model import Doc
+from models import BaseModel
 
 
 @st.cache_resource
@@ -150,9 +152,9 @@ def _get_related_url(metadata) -> Iterable[str]:
 
 def _get_query_str_filter(
     query: str,
-    project_name: str,
+    index: str,
 ) -> tuple[str, Filter]:
-    options = [{"key": "metadata.project_name", "value": project_name}]
+    options = [{"key": "metadata.index", "value": index}]
     filter = make_filter_obj(options=options)
     return query, filter
 
@@ -160,10 +162,10 @@ def _get_query_str_filter(
 def run_qa(
     llm,
     query: str,
-    project_name: str,
+    index: str,
 ) -> tuple[str, str]:
     now = time()
-    query_str, filter = _get_query_str_filter(query, project_name)
+    query_str, filter = _get_query_str_filter(query, index)
     qa = get_retrieval_qa(filter, llm)
     try:
         result = qa(query_str)
@@ -178,29 +180,30 @@ def run_qa(
 
 def run_search(
     query: str,
-    project_name: str,
-) -> Iterable[tuple[Doc, float, str]]:
-    query_str, filter = _get_query_str_filter(query, project_name)
+    index: str,
+) -> Iterable[tuple[BaseModel, float, str]]:
+    query_str, filter = _get_query_str_filter(query, index)
     qdocs = get_similay(query_str, filter)
     for qdoc, score in qdocs:
         text = qdoc.page_content
         metadata = qdoc.metadata
         # print(metadata)
-        doc = Doc(
-            project_name=project_name,
+        data = BaseModel(
+            index=index,
             id=metadata.get("id"),
             title=metadata.get("title"),
             ctime=metadata.get("ctime"),
             user=metadata.get("user"),
             url=metadata.get("url"),
+            type=metadata.get("type"),
         )
-        yield doc, score, text
+        yield data, score, text
 
 
 with st.form("my_form"):
     st.title("Document Search")
     query = st.text_input(label="query")
-    project_name = st.text_input(label="project")
+    index = st.text_input(label="index")
 
     submit_col1, submit_col2 = st.columns(2)
     searched = submit_col1.form_submit_button("Search")
@@ -209,7 +212,7 @@ with st.form("my_form"):
         st.header("Search Results")
         st.divider()
         with st.spinner("Searching..."):
-            results = run_search(query, project_name)
+            results = run_search(query, index)
             for doc, score, text in results:
                 title = doc.title
                 url = doc.url
@@ -232,7 +235,7 @@ with st.form("my_form"):
             results = run_qa(
                 LLM,
                 query,
-                project_name,
+                index,
             )
             answer, html = results
             with st.container():
@@ -249,7 +252,7 @@ with st.form("my_form"):
                 results = run_qa(
                     VICUNA_LLM,
                     query,
-                    project_name,
+                    index,
                 )
                 answer, html = results
                 with st.container():

loaders/__init__.py

@@ -0,0 +1,11 @@
+from .wikipage import WikiPageLoader
+from .github_issue import GithubIssueLoader
+
+LOADERS = {
+    "wikipage": WikiPageLoader,
+    "github_issue": GithubIssueLoader
+}
+LOADER_NAMES = tuple(LOADERS.keys())
+
+def get_loader(loader_name, **kwargs):
+    return LOADERS.get(loader_name)(**kwargs)

loaders/github_issue.py

@@ -0,0 +1,63 @@
+import json
+from dataclasses import asdict
+from pathlib import Path
+from typing import Iterator
+
+from dateutil.parser import parse
+from langchain.docstore.document import Document
+from langchain.document_loaders.base import BaseLoader
+
+from models import GithubIssue
+
+
+def date_to_int(dt_str: str) -> int:
+    dt = parse(dt_str)
+    return int(dt.timestamp())
+
+
+def get_contents(index: str, inputfile: Path) -> Iterator[tuple[GithubIssue, str]]:
+    with inputfile.open("r") as f:
+        obj = [json.loads(line) for line in f]
+    for data in obj:
+        title = data["title"]
+        body = data["body"]
+        issue = GithubIssue(
+            index=index,
+            id=data["number"],
+            title=title,
+            ctime=date_to_int(data["created_at"]),
+            user=data["user.login"],
+            url=data["html_url"],
+            labels=data["labels_"],
+        )
+        text = title
+        if body:
+            text += "\n\n" + body
+        yield issue, text
+        comments = data["comments_"]
+        for comment in comments:
+            issue = GithubIssue(
+                index=index,
+                id=comment["id"],
+                title=data["title"],
+                ctime=date_to_int(comment["created_at"]),
+                user=comment["user.login"],
+                url=comment["html_url"],
+                labels=data["labels_"],
+                type="issue_comment",
+            )
+            yield issue, comment["body"]
+
+
+class GithubIssueLoader(BaseLoader):
+    def __init__(self, index: str, inputfile: Path):
+        self.index = index
+        self.inputfile = inputfile
+
+    def lazy_load(self) -> Iterator[Document]:
+        for issue, text in get_contents(self.index, self.inputfile):
+            metadata = asdict(issue)
+            yield Document(page_content=text, metadata=metadata)
+
+    def load(self) -> list[Document]:
+        return list(self.lazy_load())

doc_loader.py → loaders/wikipage.py

@@ -1,11 +1,13 @@
-from dataclasses import asdict
 import json
+from dataclasses import asdict
+from pathlib import Path
 from typing import Iterator
+
 from dateutil.parser import parse
 from langchain.docstore.document import Document
 from langchain.document_loaders.base import BaseLoader
 
-from model import Doc
+from models import WikiPage
 
 
 def date_to_int(dt_str: str) -> int:
@@ -13,20 +15,20 @@ def date_to_int(dt_str: str) -> int:
     return int(dt.timestamp())
 
 
-def get_contents(project_name: str, filename: str) -> Iterator[tuple[Doc, str]]:
+def get_contents(index: str, inputfile: Path) -> Iterator[tuple[WikiPage, str]]:
     """filename for file with ndjson
 
         {"id": <page_id>, "title": <page title>, "content": <page body>, "ctime": ..., "user": <name>, "url": "https:..."}
         {"title": ...}
     """
-    with open(filename, "r") as f:
+    with inputfile.open("r") as f:
         obj = [json.loads(line) for line in f]
     for data in obj:
         title = data["title"]
         body = data["content"]
         ctime = date_to_int(data["ctime"]) if isinstance(data["ctime"], str) else data["ctime"]
-        doc = Doc(
-            project_name=project_name,
+        doc = WikiPage(
+            index=index,
             id=data["id"],
             title=title,
             ctime=ctime,
@@ -39,13 +41,13 @@ def get_contents(project_name: str, filename: str) -> Iterator[tuple[Doc, str]]:
         yield doc, text
 
 
-class DocLoader(BaseLoader):
-    def __init__(self, project_name: str, filename: str):
-        self.project_name = project_name
-        self.filename = filename
+class WikiPageLoader(BaseLoader):
+    def __init__(self, index: str, inputfile: Path):
+        self.index = index
+        self.inputfile = inputfile
 
     def lazy_load(self) -> Iterator[Document]:
-        for doc, text in get_contents(self.project_name, self.filename):
+        for doc, text in get_contents(self.index, self.inputfile):
             metadata = asdict(doc)
             yield Document(page_content=text, metadata=metadata)
 

model.py

@@ -1,11 +0,0 @@
-from dataclasses import dataclass
-
-
-@dataclass(frozen=True)
-class Doc:
-    project_name: str
-    id: int
-    title: str
-    ctime: int
-    user: str
-    url: str

models.py

@@ -0,0 +1,23 @@
+import dataclasses
+
+
+@dataclasses.dataclass()
+class BaseModel:
+    index: str
+    id: int
+    title: str
+    ctime: int
+    user: str
+    url: str
+    type: str
+
+
+@dataclasses.dataclass(frozen=True)
+class GithubIssue(BaseModel):
+    labels: list[str]
+    type: str = "issue"
+
+
+@dataclasses.dataclass(frozen=True)
+class WikiPage:
+    type: str = "wiki"

store.py

@@ -1,10 +1,14 @@
+import argparse
+from itertools import islice
+from pathlib import Path
+
 from tqdm import tqdm
 import torch
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import Qdrant
 
-from doc_loader import DocLoader
+from loaders import get_loader, LOADER_NAMES
 from config import DB_CONFIG
 
 
@@ -19,6 +23,16 @@ def get_text_chunk(docs):
     return texts
 
 
+def batched(iterable, *, size=100):
+    "Batch data into tuples of length n. The last batch may be shorter."
+    # batched('ABCDEFG', 3) --> ABC DEF G
+    if size < 1:
+        raise ValueError('n must be at least one')
+    it = iter(iterable)
+    while batch := tuple(islice(it, size)):
+        yield batch
+
+
 def store(texts):
     model_name = "intfloat/multilingual-e5-large"
     model_kwargs = {"device": "cuda:0" if torch.cuda.is_available() else "cpu"}
@@ -29,9 +43,9 @@ def store(texts):
         encode_kwargs=encode_kwargs,
     )
     db_url, db_api_key, db_collection_name = DB_CONFIG
-    for text in tqdm(texts):
+    for batch in tqdm(batched(texts, size=100)):
         _ = Qdrant.from_documents(
-            [text],
+            batch,
             embeddings,
             url=db_url,
             api_key=db_api_key,
@@ -39,24 +53,31 @@ def store(texts):
         )
 
 
-def main(project_name: str, path: str) -> None:
-    loader = DocLoader(project_name, path)
+def get_parser():
+    p = argparse.ArgumentParser()
+    p.add_argument("index", type=str)
+    p.add_argument("inputfile", metavar="INPUTFILE", type=argparse.FileType("rt"))
+    p.add_argument("-l", "--loader", type=str, choices=LOADER_NAMES, required=True)
+    return p
+
+
+def main():
+    """
+    $ python store.py --loader wikipage "index" "FILE_PATH"
+    $ python store.py -l wikipage wiki data/wiki.json
+    """
+    p = get_parser()
+    args = p.parse_args()
+    loader = get_loader(
+        args.loader,
+        index=args.index,
+        inputfile=Path(args.inputfile.name),
+    )
+
     docs = loader.load()
     texts = get_text_chunk(docs)
     store(texts)
 
 
 if __name__ == "__main__":
-    """
-    $ python store.py "PROJECT_NAME" "FILE_PATH"
-    $ python store.py hoge data/hoge-docs.json
-    """
-    import sys
-
-    args = sys.argv
-    if len(args) != 3:
-        print("No args, you need two args for project_name, json_file_path")
-    else:
-        project_name = args[1]
-        path = args[2]
-        main(project_name, path)
+    main()