Merge pull request #10 from ittia-research/dev

change all LLM calling to DSPy, increase citation token limit

Author: etwk

.env

@@ -1,6 +1,7 @@
 EMBEDDING_API_KEY=ollama:abc
 EMBEDDING_MODEL_DEPLOY=api
 EMBEDDING_MODEL_NAME=jina/jina-embeddings-v2-base-en
+INDEX_CHUNK_SIZES=[2048, 512, 128]
 LLM_MODEL_NAME=google/gemma-2-27b-it
 OLLAMA_BASE_URL=http://ollama:11434
 OPENAI_API_KEY=sk-proj-aaaaaaaaaaaaaaaaa
@@ -10,5 +11,4 @@ RERANK_MODEL_DEPLOY=local
 RERANK_MODEL_NAME=BAAI/bge-reranker-v2-m3
 RERANK_BASE_URL=http://xinference:9997/v1
 SEARCH_BASE_URL=https://s.jina.ai
-THREAD_BUILD_INDEX=12
-RAG_CHUNK_SIZES=[4096, 1024, 256]
+THREAD_BUILD_INDEX=12

README.md

@@ -106,6 +106,7 @@ DSPy:
 ### Reports
 - [ ] AI-generated misinformation
 ### Factcheck
+- https://www.snopes.com
 - https://www.bmi.bund.de/SharedDocs/schwerpunkte/EN/disinformation/examples-of-russian-disinformation-and-the-facts.html
 ### Resources
 #### Inference

src/llm.py

@@ -4,7 +4,7 @@
 from fastapi.responses import Response, JSONResponse, HTMLResponse, PlainTextResponse, FileResponse
 import logging
 
-import llm, utils, pipeline
+import utils, pipeline
 
 logging.basicConfig(
     level=logging.INFO,
@@ -18,7 +18,7 @@
 async def fact_check(input):
     status = 500
     logger.info(f"Fact checking: {input}")
-    statements = await run_in_threadpool(llm.get_statements, input)
+    statements = await run_in_threadpool(pipeline.get_statements, input)
     logger.info(f"statements: {statements}")
     if not statements:
         raise HTTPException(status_code=status, detail="No statements found")
@@ -29,11 +29,11 @@ async def fact_check(input):
         if not statement:
             continue
         logger.info(f"statement: {statement}")
-        keywords = await run_in_threadpool(llm.get_search_keywords, statement)
-        if not keywords:
+        query = await run_in_threadpool(pipeline.get_search_query, statement)
+        if not query:
             continue
-        logger.info(f"keywords: {keywords}")
-        search = await utils.search(keywords)
+        logger.info(f"search query: {query}")
+        search = await utils.search(query)
         if not search:
             fail_search = True
             continue

src/main.py

@@ -0,0 +1,17 @@
+import dspy
+
+from settings import settings
+
+# set DSPy default language model
+llm = dspy.OpenAI(model=settings.LLM_MODEL_NAME, api_base=f"{settings.OPENAI_BASE_URL}/", max_tokens=200, stop='\n\n')
+dspy.settings.configure(lm=llm)
+
+# LM with higher token limits
+llm_long = dspy.OpenAI(model=settings.LLM_MODEL_NAME, api_base=f"{settings.OPENAI_BASE_URL}/", max_tokens=500, stop='\n\n')
+
+from .citation import Citation
+from .ollama_embedding import OllamaEmbedding
+from .retrieve import LlamaIndexRM
+from .search_query import SearchQuery
+from .statements import Statements
+from .verdict import Verdict

src/modules/__init__.py

@@ -0,0 +1,20 @@
+import dspy
+
+# TODO: citation needs higher token limits
+class GenerateCitedParagraph(dspy.Signature):
+    """Generate a paragraph with citations."""
+    context = dspy.InputField(desc="may contain relevant facts")
+    statement = dspy.InputField()
+    verdict = dspy.InputField()
+    paragraph = dspy.OutputField(desc="includes citations")
+
+"""Generate citation from context and verdict"""
+class Citation(dspy.Module):
+    def __init__(self):
+        super().__init__()
+        self.generate_cited_paragraph = dspy.ChainOfThought(GenerateCitedParagraph)
+
+    def forward(self, statement, context, verdict):
+        citation = self.generate_cited_paragraph(context=context, statement=statement, verdict=verdict)
+        pred = dspy.Prediction(verdict=verdict, citation=citation.paragraph, context=context)
+        return pred

src/modules/citation.py

@@ -8,11 +8,9 @@
 
 from llama_index.core import (
     Document,
-    ServiceContext,
     Settings,
     StorageContext,
     VectorStoreIndex,
-    load_index_from_storage,
 )
 from llama_index.core.node_parser import HierarchicalNodeParser, get_leaf_nodes
 from llama_index.core.retrievers import AutoMergingRetriever
@@ -29,7 +27,7 @@
 jinaai_rerank.API_URL = settings.RERANK_BASE_URL + "/rerank"  # switch to on-premise
 
 # todo: high lantency between client and the ollama embedding server will slow down embedding a lot
-from ollama_embedding import OllamaEmbedding
+from . import OllamaEmbedding
 
 # todo: improve embedding performance
 if settings.EMBEDDING_MODEL_DEPLOY == "local":
@@ -132,7 +130,7 @@ def build_index(self, docs):
         if docs:
             self.index, self.storage_context = self.build_automerging_index(
                 docs,
-                chunk_sizes=settings.RAG_CHUNK_SIZES,
+                chunk_sizes=settings.INDEX_CHUNK_SIZES,
             )  # TODO: try to retrieve directly
 
     def retrieve(self, query):

src/ollama_embedding.py renamed to src/modules/ollama_embedding.py

@@ -0,0 +1,18 @@
+import dspy
+import logging
+
+"""Notes: LLM will choose a direction based on known facts"""
+class GenerateSearchEngineQuery(dspy.Signature):
+    """Write a search engine query that will help retrieve info related to the statement."""
+    statement = dspy.InputField()
+    query = dspy.OutputField()
+    
+class SearchQuery(dspy.Module):
+    def __init__(self):
+        super().__init__()
+        self.generate_query = dspy.ChainOfThought(GenerateSearchEngineQuery)
+
+    def forward(self, statement):
+        query = self.generate_query(statement=statement)
+        logging.info(f"DSPy CoT search query: {query}")
+        return query.query

src/retrieve.py renamed to src/modules/retrieve.py

@@ -0,0 +1,24 @@
+import dspy
+import logging
+from pydantic import BaseModel, Field
+from typing import List
+
+# references: https://github.com/weaviate/recipes/blob/main/integrations/llm-frameworks/dspy/4.Structured-Outputs-with-DSPy.ipynb
+class Output(BaseModel):
+    statements: List = Field(description="A list of key statements")
+    
+# TODO: test consistency especially when content contains false claims
+class GenerateStatements(dspy.Signature):
+    """Extract the original statements from given content without fact check."""
+    content: str = dspy.InputField(desc="The content to summarize")
+    output: Output = dspy.OutputField()
+
+class Statements(dspy.Module):
+    def __init__(self):
+        super().__init__()
+        self.generate_statements = dspy.TypedChainOfThought(GenerateStatements, max_retries=6)
+
+    def forward(self, content):
+        statements = self.generate_statements(content=content)
+        logging.info(f"DSPy CoT statements: {statements}")
+        return statements.output.statements

src/modules/search_query.py

@@ -1,12 +1,6 @@
 import dspy
 from dsp.utils import deduplicate
 
-from retrieve import LlamaIndexRM
-from settings import settings
-
-llm = dspy.OpenAI(model=settings.LLM_MODEL_NAME, api_base=f"{settings.OPENAI_BASE_URL}/", max_tokens=200, stop='\n\n')
-dspy.settings.configure(lm=llm)
-
 class CheckStatementFaithfulness(dspy.Signature):
     """Verify that the statement is based on the provided context."""
     context = dspy.InputField(desc="facts here are assumed to be true")
@@ -19,14 +13,6 @@ class GenerateSearchQuery(dspy.Signature):
     statement = dspy.InputField()
     query = dspy.OutputField()
 
-# TODO: citation needs higher token limits
-class GenerateCitedParagraph(dspy.Signature):
-    """Generate a paragraph with citations."""
-    context = dspy.InputField(desc="may contain relevant facts")
-    statement = dspy.InputField()
-    verdict = dspy.InputField()
-    paragraph = dspy.OutputField(desc="includes citations")
-
 """
 SimplifiedBaleen module
 Avoid unnecessary content in module cause MIPROv2 optimizer will analize modules.
@@ -39,7 +25,7 @@ class GenerateCitedParagraph(dspy.Signature):
   - remove some contexts incase token reaches to max
   - does different InputField name other than answer compateble with dspy evaluate
 """
-class ContextVerdict(dspy.Module):
+class Verdict(dspy.Module):
     def __init__(self, retrieve, passages_per_hop=3, max_hops=3):
         super().__init__()
         # self.generate_query = dspy.ChainOfThought(GenerateSearchQuery)  # IMPORTANT: solves error `list index out of range`
@@ -59,42 +45,4 @@ def forward(self, statement):
         verdict = self.generate_verdict(context=context, statement=statement)
         pred = dspy.Prediction(answer=verdict.verdict, rationale=verdict.rationale, context=context)
         return pred
-
-"""Generate citation from context and verdict"""
-class Citation(dspy.Module):
-    def __init__(self):
-        super().__init__()
-        self.generate_cited_paragraph = dspy.ChainOfThought(GenerateCitedParagraph)
-
-    def forward(self, statement, context, verdict):
-        citation = self.generate_cited_paragraph(context=context, statement=statement, verdict=verdict)
-        pred = dspy.Prediction(verdict=verdict, citation=citation.paragraph, context=context)
-        return pred
-
-"""
-Get both verdict and citation.
-
-Args:
-    retrieve: dspy.Retrieve
-"""
-class VerdictCitation():
-    def __init__(
-        self,
-        docs,
-    ):
-        self.retrieve = LlamaIndexRM(docs=docs)
-        
-        # loading compiled ContextVerdict
-        self.context_verdict = ContextVerdict(retrieve=self.retrieve)
-        self.context_verdict.load("./optimizers/verdict_MIPROv2.json")
-
-    def get(self, statement):
-        rep = self.context_verdict(statement)
-        context = rep.context
-        verdict = rep.answer
-        
-        rep = Citation()(statement=statement, context=context, verdict=verdict)
-        citation = rep.citation
-
-        return verdict, citation
-        
+ 

src/modules/statements.py

@@ -0,0 +1,2 @@
+from .common import get_search_query, get_statements, get_verdict
+from .verdict_citation import VerdictCitation

src/dspy_modules.py renamed to src/modules/verdict.py

@@ -0,0 +1,36 @@
+import logging
+import utils
+from modules import SearchQuery, Statements
+from .verdict_citation import VerdictCitation
+
+def get_statements(content):
+    """Get list of statements from a text string"""
+    try:
+        statements = Statements()(content=content)
+    except Exception as e:
+        logging.error(f"Getting statements failed: {e}")
+        statements = []
+
+    return statements
+
+def get_search_query(statement):
+    """Get search query from one statement"""
+
+    try:
+        query = SearchQuery()(statement=statement)
+    except Exception as e:
+        logging.error(f"Getting search query from statement '{statement}' failed: {e}")
+        query = ""
+
+    return query
+    
+def get_verdict(search_json, statement):
+    docs = utils.search_json_to_docs(search_json)
+    rep = VerdictCitation(docs=docs).get(statement=statement)
+
+    return {
+        "verdict": rep[0],
+        "citation": rep[1],
+        "statement": statement,
+    }
+