query_rephrase

app.py

@@ -1,156 +1,60 @@
 import gradio as gr
-import datetime
+from pipe import pipe
 import pytz
-
-sentense_transformers_model = "Alibaba-NLP/gte-multilingual-base"
-ranker_model = 'Alibaba-NLP/gte-multilingual-reranker-base'
-
-from haystack.document_stores.in_memory import InMemoryDocumentStore
-from haystack.components.retrievers.in_memory import InMemoryBM25Retriever, InMemoryEmbeddingRetriever
-from haystack.components.embedders import SentenceTransformersTextEmbedder
-from haystack.components.joiners import DocumentJoiner
-from haystack.components.rankers import TransformersSimilarityRanker
-from haystack import Pipeline,component,Document
-from haystack_integrations.components.generators.google_ai import GoogleAIGeminiChatGenerator
-from haystack.components.builders import ChatPromptBuilder
-from haystack_experimental.chat_message_stores.in_memory import InMemoryChatMessageStore
-from haystack_experimental.components.retrievers import ChatMessageRetriever
-from haystack_experimental.components.writers import ChatMessageWriter
+import datetime
 from haystack.dataclasses import ChatMessage
-from itertools import chain
-from typing import Any,List
-from haystack.core.component.types import Variadic
-from fugashi import Tagger
 
-document_store = InMemoryDocumentStore.load_from_disk(path='./document_store.json')
-print('document_store loaded' ,document_store.count_documents())
+def log_docs(docs):
+    for doc in docs :
+        title = doc.meta['title']
+        link = doc.meta['link']
+        print('',title,link,doc.meta['type'],doc.score)
+def get_unique_docs(docs):
+    source_ids = set([doc.meta['source_id'] for doc in docs])
+    _docs = sorted([[doc for doc in docs if doc.meta['source_id']==source_id][0] for source_id in source_ids],key=lambda x:x.score,reverse=True)
+    return _docs
+     
+def run(pipe,q):
+    now = datetime.datetime.now(pytz.timezone('Asia/Tokyo'))
+    print('\nq:',q,now)
+    if not q :
+        return {'reply':'','sources':''}
+    
+    result = pipe.run({
+        "query_rephrase_prompt_builder":{"query":q},
+        "prompt_builder": { "query": q},
+        "memory_joiner": {"values": [ChatMessage.from_user(q)]},
+    },include_outputs_from=["query_rephrase_llm","llm",'ranker','bm25_retriever','embedding_retriever'])
 
-@component
-class ListJoiner:
-    def __init__(self, _type: Any):
-        component.set_output_types(self, values=_type)
-    def run(self, values: Variadic[Any]):
-        result = list(chain(*values))
-        return {"values": result}
-
-tagger = Tagger()
-def gen_wakachi(text):
-    words = tagger(text)
-    return ' '.join([word.surface for word in words])
-
-class Niwa_rag :
-    def __init__(self):
-        self.createPipe()
-    def createPipe(self):
-        user_message_template = """
-            チャット履歴と提供された資料に基づいて、質問に答えてください。
-            資料はチャット履歴の一部ではないことに注意してください。
-            質問が資料から回答できない場合は、その旨を述べてください。
-            
-            チャット履歴:
-            {% for memory in memories %}
-                {{ memory.content }}
-            {% endfor %}
-        
-            資料:
-            {% for document in documents %}
-                {{ document.content }}
-            {% endfor %}
-        
-            質問: {{query}}
-            回答:
-        """
-        system_message = ChatMessage.from_system("あなたは、提供された資料とチャット履歴を使用して人間を支援するAIアシスタントです")
-        user_message = ChatMessage.from_user(user_message_template)
-        messages = [system_message, user_message]
+    query_rephrase = result['query_rephrase_llm']['replies'][0]
+    print('query_rephrase:',query_rephrase)
+    
+    for retriever in ['bm25_retriever','embedding_retriever'] :
+        print(retriever)
+        docs = result[retriever]['documents']
+        log_docs(docs)
         
-        text_embedder = SentenceTransformersTextEmbedder(model=sentense_transformers_model,trust_remote_code=True)
-        embedding_retriever = InMemoryEmbeddingRetriever(document_store)
-        bm25_retriever = InMemoryBM25Retriever(document_store)
-        document_joiner = DocumentJoiner()
-        ranker = TransformersSimilarityRanker(model=ranker_model,meta_fields_to_embed=['title'],model_kwargs={'trust_remote_code':True})
-        prompt_builder = ChatPromptBuilder(template=messages,variables=["query", "documents", "memories"], required_variables=["query", "documents", "memories"])
-        gemini = GoogleAIGeminiChatGenerator(model="models/gemini-1.5-flash")
-        
-        memory_store = InMemoryChatMessageStore()
-        memory_joiner = ListJoiner(List[ChatMessage])
-        memory_retriever = ChatMessageRetriever(memory_store)
-        memory_writer = ChatMessageWriter(memory_store)
-        
-        pipe = Pipeline()
-        pipe.add_component("text_embedder", text_embedder)
-        pipe.add_component("embedding_retriever", embedding_retriever)
-        pipe.add_component("bm25_retriever", bm25_retriever)
-        pipe.add_component("document_joiner", document_joiner)
-        pipe.add_component("ranker", ranker)
-        pipe.add_component("prompt_builder", prompt_builder)
-        pipe.add_component("llm", gemini)
-        pipe.add_component("memory_retriever", memory_retriever)
-        pipe.add_component("memory_writer", memory_writer)
-        pipe.add_component("memory_joiner", memory_joiner)
-        
-        pipe.connect("text_embedder", "embedding_retriever")
-        pipe.connect("bm25_retriever", "document_joiner")
-        pipe.connect("embedding_retriever", "document_joiner")
-        pipe.connect("document_joiner", "ranker")
-        pipe.connect("ranker.documents", "prompt_builder.documents")
-        pipe.connect("prompt_builder.prompt", "llm.messages")
-        pipe.connect("llm.replies", "memory_joiner")
-        pipe.connect("memory_joiner", "memory_writer")
-        pipe.connect("memory_retriever", "prompt_builder.memories")
-        
-        self.pipe = pipe
-    def run(self,q):
-        now = datetime.datetime.now(pytz.timezone('Asia/Tokyo'))
-        print('\nq:',q,now)
-        if not q :
-            return {'reply':'','sources':''}
-        
-        result = self.pipe.run({
-            "text_embedder": {"text": q},
-            "bm25_retriever": {"query": gen_wakachi(q)},
-            "ranker": {"query": q},
-            "prompt_builder": { "query": q},
-            "memory_joiner": {"values": [ChatMessage.from_user(q)]},
-        },include_outputs_from=["llm",'ranker','bm25_retriever','embedding_retriever'])
-        
-        def log_document(doc):
-            title = doc.meta['title']
-            link = doc.meta['link']
-            print('',title,link,doc.meta['type'],doc.score)
-            
-        for retriever in ['bm25_retriever','embedding_retriever'] :
-            print(retriever)
-            docs = result[retriever]['documents']
-            for doc in docs :
-                log_document(doc)
-            
-        reply = result['llm']['replies'][0]
-        docs = result['ranker']['documents']
-        print('reply:',reply)
-        print('ranker')
-        def get_unique_docs(docs):
-            source_ids = set([doc.meta['source_id'] for doc in docs])
-            _docs = sorted([[doc for doc in docs if doc.meta['source_id']==source_id][0] for source_id in source_ids],key=lambda x:x.score,reverse=True)
-            return _docs
-        html = '<div class="ref-title">参考記事</div><div class="ref">'
-        for doc in get_unique_docs(docs) :
-            log_document(doc)
-            title = doc.meta['title']
-            link = doc.meta['link']
-            row = f'<div><a class="link" href="{link}" target="_blank">{title}</a></div>'
-            html += row
-        html += '</div>'
-        return {'reply':reply.content,'sources':html}
-
-rag = Niwa_rag()
+    reply = result['llm']['replies'][0]
+    docs = result['ranker']['documents']
+    print('reply:',reply.content)
+    print('ranker')
+    log_docs(docs)
+    
+    html = '<div class="ref-title">参考記事</div><div class="ref">'
+    for doc in get_unique_docs(docs) :
+        title = doc.meta['title']
+        link = doc.meta['link']
+        row = f'<div><a class="link" href="{link}" target="_blank">{title}</a></div>'
+        html += row
+    html += '</div>'
+    return {'reply':reply.content,'sources':html}
 
-def fn(q,history):
-    result = rag.run(q)
+def rag(q,history):
+    result = run(pipe,q)
     return result['reply'] + result['sources']
 
 app = gr.ChatInterface(
-    fn, 
+    rag, 
     type="messages",
     title='庭ファン　Chatbot',
     textbox=gr.Textbox(placeholder='質問を記入して下さい',submit_btn=True),

pipe.py

@@ -0,0 +1,132 @@
+sentense_transformers_model = "Alibaba-NLP/gte-multilingual-base"
+ranker_model = 'Alibaba-NLP/gte-multilingual-reranker-base'
+gemini_model = 'models/gemini-1.5-flash'
+
+
+from haystack.document_stores.in_memory import InMemoryDocumentStore
+from haystack.components.retrievers.in_memory import InMemoryBM25Retriever, InMemoryEmbeddingRetriever
+from haystack.components.embedders import SentenceTransformersTextEmbedder
+from haystack.components.joiners import DocumentJoiner
+from haystack.components.rankers import TransformersSimilarityRanker
+from haystack import Pipeline,component,Document
+from haystack_integrations.components.generators.google_ai import GoogleAIGeminiChatGenerator,GoogleAIGeminiGenerator
+from haystack.components.builders import ChatPromptBuilder,PromptBuilder
+from haystack_experimental.chat_message_stores.in_memory import InMemoryChatMessageStore
+from haystack_experimental.components.retrievers import ChatMessageRetriever
+from haystack_experimental.components.writers import ChatMessageWriter
+from haystack.dataclasses import ChatMessage
+from itertools import chain
+from typing import Any,List
+from haystack.core.component.types import Variadic
+from haystack.components.converters import OutputAdapter
+from fugashi import Tagger
+
+tagger = Tagger()
+def gen_wakachi(text):
+    words = tagger(text)
+    return ' '.join([word.surface for word in words])
+
+@component
+class ListJoiner:
+    def __init__(self, _type: Any):
+        component.set_output_types(self, values=_type)
+    def run(self, values: Variadic[Any]):
+        result = list(chain(*values))
+        return {"values": result}
+
+@component 
+class WakachiAdapter:
+    @component.output_types(wakachi=str)
+    def run(self, query:str):
+        return {'wakachi':gen_wakachi(query)}
+
+document_store = InMemoryDocumentStore.load_from_disk(path='./document_store.json')
+print('document_store loaded' ,document_store.count_documents())
+
+user_message_template = """
+    チャット履歴と提供された資料に基づいて、質問に答えてください。
+    資料はチャット履歴の一部ではないことに注意してください。
+    質問が資料から回答できない場合は、その旨を述べてください。
+    
+    チャット履歴:
+    {% for memory in memories %}
+        {{ memory.content }}
+    {% endfor %}
+
+    資料:
+    {% for document in documents %}
+        {{ document.content }}
+    {% endfor %}
+
+    質問: {{query}}
+    回答:
+"""
+query_rephrase_template = """
+    意味とキーワードをそのまま維持しながら、検索用の質問を書き直してください。
+    チャット履歴が空の場合は、クエリを変更しないでください。
+    チャット履歴は必要な場合にのみ使用し、独自の知識でクエリを拡張することは避けてください。
+    変更の必要がない場合は、現在の質問をそのまま出力して下さい。
+
+    チャット履歴 :
+    {% for memory in memories %}
+        {{ memory.text }}
+    {% endfor %}
+
+    ユーザーの質問 : {{query}}
+    書き換えられた質問 :
+"""
+
+system_message = ChatMessage.from_system("あなたは、提供された資料とチャット履歴を使用して人間を支援するAIアシスタントです")
+user_message = ChatMessage.from_user(user_message_template)
+messages = [system_message, user_message]
+
+query_rephrase_prompt_builder = PromptBuilder(query_rephrase_template)
+query_rephrase_llm = GoogleAIGeminiGenerator(model=gemini_model)
+list_to_str_adapter = OutputAdapter(template="{{ replies[0] }}", output_type=str)
+wakachi_adapter = WakachiAdapter()
+
+text_embedder = SentenceTransformersTextEmbedder(model=sentense_transformers_model,trust_remote_code=True)
+embedding_retriever = InMemoryEmbeddingRetriever(document_store)
+bm25_retriever = InMemoryBM25Retriever(document_store)
+document_joiner = DocumentJoiner()
+ranker = TransformersSimilarityRanker(model=ranker_model,meta_fields_to_embed=['title'],model_kwargs={'trust_remote_code':True})
+prompt_builder = ChatPromptBuilder(template=messages,variables=["query", "documents", "memories"], required_variables=["query", "documents", "memories"])
+gemini = GoogleAIGeminiChatGenerator(model=gemini_model)
+
+memory_store = InMemoryChatMessageStore()
+memory_joiner = ListJoiner(List[ChatMessage])
+memory_retriever = ChatMessageRetriever(memory_store)
+memory_writer = ChatMessageWriter(memory_store)
+
+pipe = Pipeline()
+pipe.add_component("query_rephrase_prompt_builder", query_rephrase_prompt_builder)
+pipe.add_component("query_rephrase_llm", query_rephrase_llm)
+pipe.add_component("list_to_str_adapter", list_to_str_adapter)
+pipe.add_component("wakachi_adapter", wakachi_adapter)
+pipe.add_component("text_embedder", text_embedder)
+pipe.add_component("embedding_retriever", embedding_retriever)
+pipe.add_component("bm25_retriever", bm25_retriever)
+pipe.add_component("document_joiner", document_joiner)
+pipe.add_component("ranker", ranker)
+pipe.add_component("prompt_builder", prompt_builder)
+pipe.add_component("llm", gemini)
+pipe.add_component("memory_retriever", memory_retriever)
+pipe.add_component("memory_joiner", memory_joiner)
+pipe.add_component("memory_writer", memory_writer)
+
+pipe.connect("memory_retriever", "query_rephrase_prompt_builder.memories")
+pipe.connect("query_rephrase_prompt_builder.prompt", "query_rephrase_llm")
+pipe.connect("query_rephrase_llm.replies", "list_to_str_adapter")
+pipe.connect("list_to_str_adapter", "text_embedder.text")
+pipe.connect("list_to_str_adapter", "wakachi_adapter.query")
+pipe.connect("wakachi_adapter.wakachi", "bm25_retriever.query")
+pipe.connect("text_embedder", "embedding_retriever")
+pipe.connect("embedding_retriever", "document_joiner")
+pipe.connect("bm25_retriever", "document_joiner")
+pipe.connect("document_joiner", "ranker")
+pipe.connect("list_to_str_adapter", "ranker.query")
+pipe.connect("ranker.documents", "prompt_builder.documents")
+pipe.connect("memory_retriever", "prompt_builder.memories")
+pipe.connect("prompt_builder.prompt", "llm.messages")
+pipe.connect("llm.replies", "memory_joiner")
+pipe.connect("memory_joiner", "memory_writer")