Aman/

selfrag-zh_baichuan2_7b_chat

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>