add similarity check

Author: tmori

GenerativeAgentForDoq.ipynb

@@ -74,6 +74,7 @@ pip3 install pypdf
 pip3 install langchain
 pip3 install unstructured
 pip3 install tabulate
+pip3 install scikit-learn
 ```
 
 Please set the OpenAPI API key as an environment variable.

README.md

@@ -0,0 +1,37 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+
+import openai
+import tiktoken
+from tiktoken.core import Encoding
+from openai.embeddings_utils import cosine_similarity
+
+llm_model = "gpt-4-0613"
+embedding_model = "text-embedding-ada-002"
+
+# Embedding
+def get_embedding(text_input: str):
+    global embedding_model
+    
+    # ベクトル変換
+    response  = openai.Embedding.create(
+                    input = text_input.replace("\n", " "),   # 入力文章
+                    model = embedding_model,        # GPTモデル
+                 )
+    
+    # 出力結果取得
+    embeddings = response['data'][0]['embedding']
+    
+    return embeddings
+
+def get_score(text1: str, text2: str):
+    vec1 = get_embedding(text1)
+    vec2 = get_embedding(text2)
+    result = cosine_similarity(vec1, vec2)
+    return result
+
+
+def get_tokenlen(data: str):
+    encoding: Encoding = tiktoken.encoding_for_model(llm_model)
+    tokens = encoding.encode(data)
+    return len(tokens)

data_model/openai_libs.py

@@ -0,0 +1,64 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+
+from data_model_accessor import DataModelAccessor
+import json
+from openai_libs import get_score, get_tokenlen
+
+class SimilarityExtractor:
+    def __init__(self, accessor: DataModelAccessor, maxtoken_num: int):
+        self.maxtoken_num = maxtoken_num
+        self.accessor = accessor
+
+    def get_filelist(self, query: str):
+        scores = self._calc_scores(query, accessor.get_filelist())
+        result = []
+        token_sum = 0
+        for entry in scores:
+            if token_sum + entry["tokens"] > self.maxtoken_num:
+                break
+            result.append(entry["file"])
+            token_sum += entry["tokens"]
+        return result
+
+    def _calc_scores(self, query: str, filelist: list):
+        scores = []
+        for entry in filelist:
+            #print("file:", entry)
+            json_data = self.accessor.get_data_model(entry).get_json_data()
+            json_str = json.dumps(json_data)
+            score = get_score(query, json_str)
+            tokens = get_tokenlen(json_str)
+            scores.append({
+                "file": entry,
+                "tokens": tokens,
+                "score": score
+            })
+        scores.sort(key=lambda x: x["score"], reverse=True)
+        return scores
+
+
+    def extract(self, head_name: str, filelists: list):
+        models = self.accessor.get_json_models(filelists)
+        data = {
+            head_name: models
+        }
+        return data
+
+
+if __name__ == "__main__":
+    import sys
+    import json
+    if len(sys.argv) != 3:
+        print("Usage: <query> <dir>")
+        sys.exit(1)
+    query = sys.argv[1]
+    dir = sys.argv[2]
+
+    accessor = DataModelAccessor(dir)
+    extractor = SimilarityExtractor(accessor, 1024)
+
+    filelist = extractor.get_filelist(query)
+    data = extractor.extract("knowledges", filelist)
+    data_str = json.dumps(data, indent=4, ensure_ascii=False)
+    print(data_str)

data_model/similarity_extractor.py

@@ -44,6 +44,7 @@ do
             rm -rf test/*.json
             echo "INFO: CRITICAL THINKING"
             python3 critical_thinking.py  "$query" ${background_file}
+            python3 data_model/reflection_data_persistentor.py tmpdir ./test/result/critical_thinking.json
             echo "INFO: GETTING DOCUMENTS"
             get_docs "${query}" critical_thinking.json
             documents=`cat tmp.list`
@@ -88,6 +89,8 @@ do
         else
             python3 reflection.py "$query" ../documents/document.list "./test/result/critical_thinking.json" ${background_file} "./prompt_templates/ptemplate_reflection.txt"
         fi
+        python3 data_model/reflection_data_persistentor.py tmpdir ./test/result/reflection.json
+
         cp ./test/result/reflection.json ./test/result/prev_reflection.json
         python3 reflection.py "$query" ../documents/document.list "./test/result/reflection.json" ${background_file} "./prompt_templates/ptemplate_reflection_addterms.txt"
         mv ./test/result/reflection.json ./test/result/next_reflection.json