---
license: mit
tags:
- text generation
- RAG
- baichuan2
---

This model is a 7B Chinese version of [Self-RAG](https://huggingface.co/selfrag/selfrag_llama2_7b).

It is trained on Baichuan2-7B-Chat with a sample of [belle](https://github.com/LianjiaTech/BELLE) sft data, acompanying with interleaving passages from zhwiki. The reflection tokens are aligned with the original verison (in English), so the usage is the same. Hope you enjoy.


### Data
The data used to train the model is also available ([FINAL_OUTPUT_4w.jsonl](https://huggingface.co/Aman/selfrag-zh_baichuan2_7b_chat/blob/main/FINAL_OUTPUT_4w.jsonl)), which is constructed using [Belle](https://github.com/LianjiaTech/BELLE/tree/main/data/1.5M) SFT data and Wikipedia Chinese docs.

### Usage

#### The Critic Model
The critic model is released at the `critic/` folder. However, due to the quantity and quality of the critic data, there is still a distance from a perfect performance.

#### The Generator
I found some output errors while adopting vllm to accelerate the generation process and not sure whether it is due to some precision issues.
This may be owing to the implementation of vllm. Thus, I use the original generate method of transformers.
```
import os, torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(YOUR_TOKENIZER_PATH)
model = AutoModelForCausalLM.from_pretrained(
        YOUR_MODEL_PATH,
        torch_dtype=torch.bfloat16,
        device_map="cuda",
    )

### set your retriever if necessary
retriever = setup_retriever(YOUR_RETRIEVER_PATH)


def format_prompt(input, paragraph=None):
    prompt = "### Instruction:\n{0}\n\n### Response:".format(input)
    if paragraph is not None:
        prompt += "[Retrieval]<paragraph>{0}</paragraph>".format(paragraph)
    return prompt


while True:
    query = input("[Human]: ")
    prompt = format_prompt(query)
    sequences = model.generate(
        **tokenizer(prompt, return_tensors='pt').to(model.device),
        do_sample=False,
        num_beams=5,
        # top_k=10,
        # top_p=0.8,
        temperature=0.9,
        num_return_sequences=1,
        eos_token_id=tokenizer.eos_token_id,
        max_new_tokens=1024,
        min_new_tokens=1,
        repetition_penalty=1.5,
    )
    for seq in sequences:
        print(f"[Model]: {tokenizer.decode(seq, skip_special_tokens=False)}")
        print("-"*50)
    print("="*50)

# query_1 = "你好呀"
# Model prediction: [No Retrieval] 你好！有什么我可以帮你解答的问题吗？ [Utility:5] </s>
# query_2 = "故宫三大殿是哪些？"
# Model prediction: [Retrieval] <paragraph> ... (this query requires factual grounding, call a retriever) </paragraph> [Relevant] 太和殿、中和殿、保和殿 [Utility:5] </s>
```