pip install llama_cpp_python
pip install huggingface-hub

huggingface-cli download m4r1/selfrag_llama2_7b-GGUF selfrag_llama2_7b.q4_k_m.gguf --local-dir ./model --local-dir-use-symlinks False

from llama_cpp import Llama
_MODEL_KWARGS = {"logits_all": True, "n_ctx": 2048, "n_gpu_layers":200}
_GENERATE_KWARGS = {"temperature": 0.0,"top_p": 1.0,"max_tokens": 1024,"logprobs": 1000}

#模型
llm=Llama(model_path="./model/selfrag_llama2_7b.q4_k_m.gguf",**_MODEL_KWARGS)

#格式化Prompt，注意必须按照此格式输入问题和关联知识
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 = "写一首歌颂母爱的小诗"
query_2 = "能否介绍下字节跳动的AI平台Coze？"
queries = [query_1, query_2]

#分别测试，并打印出结果(response); details用来查看更详细的tokens输出细节
for query in queries:
    pred = llm(format_prompt(query),**_GENERATE_KWARGS)
    print("\nResponse: {0}".format(pred["choices"][0]["text"]))
    print(pred["choices"][0])

paragraph="""Coze是字节跳动的大模型应用一站式开发平台。"""
...
#注意此处把输入参数paragraph默认成上面的知识
def format_prompt(input, paragraph=paragraph):
...

#相关度计算：[Relevant]输出的概率占比
def _relevance_score(pred_log_probs: Dict[str, float]) -> float:
    rel_prob = np.exp(float(pred_log_probs["[Relevant]"]))
    irel_prob = np.exp(float(pred_log_probs["[Irrelevant]"]))
    return rel_prob / (rel_prob + irel_prob)

import os
from llama_index.llms.llama_cpp import LlamaCPP
from llama_index.core import Document, VectorStoreIndex
from llama_index.core.retrievers import VectorIndexRetriever
from pathlib import Path

#导入SelfRAGQueryEngine引擎
from selfrag_queryengine import SelfRAGQueryEngine

#模型参数，注意打开logits_all参数
_MODEL_KWARGS = {"logits_all": True, "n_ctx": 2048, "n_gpu_layers": -1}
_GENERATE_KWARGS = {"temperature": 0.0,"top_p": 1.0,"max_tokens": 1000,"logprobs": 32016,}

# 下载的selfrag_llama2_7b模型的保存目录
download_dir = "../../model"

# 创建测试文档，此处直接用documents构建，方便观察
documents = [
    Document(text="Xiaomi 14 is the latest smartphone released by Xiaomi. It adopts a new design concept, the body is lighter and thinner, equipped with the latest processor, and the performance is more powerful."),
    """此处省略了更多的Document对象"""
]

# 创建retriever
index = VectorStoreIndex.from_documents(documents)
retriever = VectorIndexRetriever(index=index,similarity_top_k=5)

# 创建llm（采用llama_cpp）
model_path = Path(download_dir) / "selfrag_llama2_7b.q4_k_m.gguf"
llm = LlamaCPP(model_path=str(model_path), model_kwargs=_MODEL_KWARGS, generate_kwargs=_GENERATE_KWARGS)

# 构造查询引擎，传入llm与retriever
query_engine = SelfRAGQueryEngine(llm, retriever)

# 查询一：无需检索的创作问题
response = query_engine.query("write a poem about beautiful sunset")

# 查询二：需要检索的事实性问题
response = query_engine.query("Tell me some truth about xiaomi 14 phone, especially about its battery and camera?")

"""此处省略初始化代码"""def query(self, query_str: str) -> str:
            #调用模型生成
            response = self.llm.complete(_format_prompt(query_str))
            answer = response.text

            #如果模型反馈需要检索
            if "[Retrieval]" in answer:

                #检索多个相关知识
                print_text("需要检索知识，开始检索...\n", color="blue")
                documents = self.retriever.retrieve(query_str)
                print_text(f"共检索到 {len(documents)} 个相关知识\n", color="blue")

                #用检索到的多个知识组装多个Prompt
                paragraphs = [
                    _format_prompt(query_str, document.node.text) for document in documents
                ]
               
                #重新生成并评估每个结果的评分
                print_text("=====开始：重新生成并评估====\n", color="blue")
                llm_response_per_paragraph,paragraphs_final_score = \
                    self._regen_then_eval(paragraphs)
                print_text("===结束：重新生成并评估====\n", color="blue")

                #选择评分最高的答案
                best_paragraph_id = max(
                    paragraphs_final_score, key=paragraphs_final_score.get
                )
                answer = llm_response_per_paragraph[best_paragraph_id]
                print_text(f"已选择最佳答案: {answer}\n", color="blue")

            else:
                print_text("无需检索知识，直接输出答案\n",color="green")

            #输出结果，此处需要去除答案中的特殊token
            answer = _postprocess_answer(answer)
            print_text(f"最终答案: {answer}\n", color="green")
            return str(answer)

def _regen_then_eval(self, paragraphs: List[str]) ->Tuple[Dict[int,str],Dict[int,float]]:
            """
            运行评判模块，对给定的段落进行生成，并评分。
            参数：
            paragraphs (List[str]): 包含要评分的段落的列表。
            返回：
            Tuple[Dict[int,str],Dict[int,float]]: 包含生成的结果和评分。
            """
            paragraphs_final_score = {}
            llm_response_text = {}

            for p_idx, paragraph in enumerate(paragraphs):
                #循环生成多个响应
                response = self.llm.complete(paragraph)
                pred = response.raw
                llm_response_text[p_idx] = response.text

                #从raw信息中取得tokens概率相关信息
                #top_logprobs保存了每个位置上输出每个token的概率
                logprobs = pred["choices"][0]["logprobs"]
                pred_log_probs = logprobs["top_logprobs"]

                # 计算isRel分数，相关度为第一个token，直接传入0
                isRel_score = _relevance_score(pred_log_probs[0])

                # 计算isSup分数
                isSup_score = _supported_score(logprobs["tokens"], pred_log_probs)

                # 计算isUse分数
                isUse_score = _useful_score(logprobs["tokens"], pred_log_probs)

                #最终得分
                paragraphs_final_score[p_idx] = (
                    isRel_score + isSup_score + 0.5 * isUse_score
                )

                print_text(
                    f"输入: {paragraph}\n响应: {llm_response_text[p_idx]}\n评分: {paragraphs_final_score[p_idx]}\n",
                    color="blue",
                 )
                print_text(
                    f"已完成 {p_idx + 1}/{len(paragraphs)} 段落\n\n", color="blue"
                )

            return llm_response_text, paragraphs_final_score