揭秘Self-RAG：引领大型语言模型生成质量的新潮流！

from vllm import LLM, SamplingParams
model = LLM("selfrag/selfrag_llama2_7b", download_dir="/gscratch/h2lab/akari/model_cache", dtype="half")sampling_params = SamplingParams(temperature=0.0, top_p=1.0, max_tokens=100, skip_special_tokens=False)
def format_prompt(input, paragraph=None):  prompt = "### Instruction:\n{0}\n\n### Response:\n".format(input)  if paragraph is not None:    prompt += "[Retrieval]<paragraph>{0}</paragraph>".format(paragraph)  return prompt
query_1 = "Leave odd one out: twitter, instagram, whatsapp."query_2 = "Can you tell me the difference between llamas and alpacas?"queries = [query_1, query_2]
# for a query that doesn't require retrievalpreds = model.generate([format_prompt(query) for query in queries], sampling_params)for pred in preds:  print("Model prediction: {0}".format(pred.outputs[0].text))

def retrieve(state):    """    Retrieve documents
    Args:        state (dict): The current graph state
    Returns:        state (dict): New key added to state, documents, that contains retrieved documents    """    print("---RETRIEVE---")    question = state["question"]
    # Retrieval    documents = retriever.get_relevant_documents(question)    return {"documents": documents, "question": question}

def generate(state):    """    Generate answer
    Args:        state (dict): The current graph state
    Returns:        state (dict): New key added to state, generation, that contains LLM generation    """    print("---GENERATE---")    question = state["question"]    documents = state["documents"]
    # RAG generation    generation = rag_chain.invoke({"context": documents, "question": question})    return {"documents": documents, "question": question, "generation": generation}

def grade_documents(state):    """    Determines whether the retrieved documents are relevant to the question.
    Args:        state (dict): The current graph state
    Returns:        state (dict): Updates documents key with only filtered relevant documents    """
    print("---CHECK DOCUMENT RELEVANCE TO QUESTION---")    question = state["question"]    documents = state["documents"]
    # Score each doc    filtered_docs = []    for d in documents:        score = retrieval_grader.invoke(            {"question": question, "document": d.page_content}        )        grade = score.binary_score        if grade == "yes":            print("---GRADE: DOCUMENT RELEVANT---")            filtered_docs.append(d)        else:            print("---GRADE: DOCUMENT NOT RELEVANT---")            continue    return {"documents": filtered_docs, "question": question}

def transform_query(state):    """    Transform the query to produce a better question.
    Args:        state (dict): The current graph state
    Returns:        state (dict): Updates question key with a re-phrased question    """
    print("---TRANSFORM QUERY---")    question = state["question"]    documents = state["documents"]
    # Re-write question    better_question = question_rewriter.invoke({"question": question})    return {"documents": documents, "question": better_question}

### Edges

def decide_to_generate(state):    """    Determines whether to generate an answer, or re-generate a question.
    Args:        state (dict): The current graph state
    Returns:        str: Binary decision for next node to call    """
    print("---ASSESS GRADED DOCUMENTS---")    state["question"]    filtered_documents = state["documents"]
    if not filtered_documents:        # All documents have been filtered check_relevance        # We will re-generate a new query        print(            "---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY---"        )        return "transform_query"    else:        # We have relevant documents, so generate answer        print("---DECISION: GENERATE---")        return "generate"

def grade_generation_v_documents_and_question(state):    """    Determines whether the generation is grounded in the document and answers question.
    Args:        state (dict): The current graph state
    Returns:        str: Decision for next node to call    """
    print("---CHECK HALLUCINATIONS---")    question = state["question"]    documents = state["documents"]    generation = state["generation"]
    score = hallucination_grader.invoke(        {"documents": documents, "generation": generation}    )    grade = score.binary_score
    # Check hallucination    if grade == "yes":        print("---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---")        # Check question-answering        print("---GRADE GENERATION vs QUESTION---")        score = answer_grader.invoke({"question": question, "generation": generation})        grade = score.binary_score        if grade == "yes":            print("---DECISION: GENERATION ADDRESSES QUESTION---")            return "useful"        else:            print("---DECISION: GENERATION DOES NOT ADDRESS QUESTION---")            return "not useful"    else:        pprint("---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---")        return "not supported"