超GPT3.5性能，无穷长文本，超强RAG三件套，MiniCPM3-4B模型分享 - ToB企服应用市场:ToB评测及商务社交产业平台

import torch
from modelscope import AutoModelForCausalLM, AutoModel, AutoTokenizer, snapshot_download
from transformers import AutoModelForSequenceClassification
from peft import PeftModel
import torch.nn.functional as F
import numpy as np
def MiniCPM3_4B_inference(message, model_path="OpenBMB/MiniCPM3-4B", device="cuda"):
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, device_map=device, trust_remote_code=True)
messages = [{"role": "user", "content": message}]
model_inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(device)
model_outputs = model.generate(
model_inputs,
max_new_tokens=1024,
top_p=0.7,
temperature=0.7,
repetition_penalty=1.02
)
output_token_ids = [model_outputs[i][len(model_inputs[i]):] for i in range(len(model_inputs))]
responses = tokenizer.batch_decode(output_token_ids, skip_special_tokens=True)[0]
return responses
def MiniCPM3_RAG_LoRA_inference(instruction, passages_list, base_model_dir="OpenBMB/MiniCPM3-4B", lora_model_dir="OpenBMB/MiniCPM3-RAG-LoRA"):
base_model_dir = snapshot_download(base_model_dir)
lora_model_dir = snapshot_download(lora_model_dir)
model = AutoModelForCausalLM.from_pretrained(base_model_dir, device_map="auto", torch_dtype=torch.bfloat16).eval()
tokenizer = AutoTokenizer.from_pretrained(lora_model_dir)
model = PeftModel.from_pretrained(model, lora_model_dir)
passages = '\n'.join(passages_list)
input_text = 'Background:\n' + passages + '\n\n' + instruction
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": input_text},
]
prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
outputs = model.chat(tokenizer, prompt, temperature=0.8, top_p=0.8)
return outputs[0]
def MiniCPM_Embedding_inference(queries, passages, model_name="OpenBMB/MiniCPM-Embedding", device="cuda"):
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True, attn_implementation="flash_attention_2", torch_dtype=torch.float16).to(device)
model.eval()
def weighted_mean_pooling(hidden, attention_mask):
attention_mask_ = attention_mask * attention_mask.cumsum(dim=1)
s = torch.sum(hidden * attention_mask_.unsqueeze(-1).float(), dim=1)
d = attention_mask_.sum(dim=1, keepdim=True).float()
reps = s / d
return reps
@torch.no_grad()
def encode(input_texts):
batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt', return_attention_mask=True).to(device)
outputs = model(**batch_dict)
attention_mask = batch_dict["attention_mask"]
hidden = outputs.last_hidden_state
reps = weighted_mean_pooling(hidden, attention_mask)
embeddings = F.normalize(reps, p=2, dim=1).detach().cpu().numpy()
return embeddings
INSTRUCTION = "Query: "
queries = [INSTRUCTION + query for query in queries]
embeddings_query = encode(queries)
embeddings_doc = encode(passages)
scores = (embeddings_query @ embeddings_doc.T)
return scores.tolist()
def MiniCPM_Reranker_rerank(queries, passages, model_name='OpenBMB/MiniCPM-Reranker', device="cuda", max_len_q=512, max_len_d=512):
model_name = snapshot_download(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.padding_side = "right"
model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True, attn_implementation="flash_attention_2", torch_dtype=torch.float16).to(device)
model.eval()
def tokenize_our(query, doc):
input_id_query = tokenizer.encode(query, add_special_tokens=False, max_length=max_len_q, truncation=True)
input_id_doc = tokenizer.encode(doc, add_special_tokens=False, max_length=max_len_d, truncation=True)
pad_input = {"input_ids": [tokenizer.bos_token_id] + input_id_query + [tokenizer.eos_token_id] + input_id_doc}
return tokenizer.pad(
pad_input,
padding="max_length",
max_length=max_len_q + max_len_d + 2,
return_tensors="pt",
)
@torch.no_grad()
def rerank(input_query, input_docs):
tokenized_inputs = [tokenize_our(input_query, input_doc).to(device) for input_doc in input_docs]
input_ids = {
"input_ids": [tokenized_input["input_ids"] for tokenized_input in tokenized_inputs],
"attention_mask": [tokenized_input["attention_mask"] for tokenized_input in tokenized_inputs]
}
for k in input_ids:
input_ids[k] = torch.stack(input_ids[k]).to(device)
outputs = model(**input_ids)
score = outputs.logits
return score.float().detach().cpu().numpy()
INSTRUCTION = "Query: "
queries = [INSTRUCTION + query for query in queries]
scores = [rerank(query, docs) for query, docs in zip(queries, passages)]
return np.array(scores)
def main():
# Example use cases
response_4B = MiniCPM3_4B_inference("推荐5个北京的景点。")
print(f"MiniCPM3-4B Response: {response_4B}")
instruction = "Q: What is the name of the lead character in the novel 'The Silent Watcher'?\nA:"
passages_list = [
"In the novel 'The Silent Watcher,' the lead character is named Alex Carter. Alex is a private detective who uncovers a series of mysterious events in a small town.",
"Set in a quiet town, 'The Silent Watcher' follows Alex Carter, a former police officer turned private investigator, as he unravels the town's dark secrets.",
"'The Silent Watcher' revolves around Alex Carter's journey as he confronts his past while solving complex cases in his hometown."
]
response_RAG_LoRA = MiniCPM3_RAG_LoRA_inference(instruction, passages_list)
print(f"MiniCPM3-RAG-LoRA Response: {response_RAG_LoRA}")
queries = ["China capital?"]
passages = ["beijing", "shanghai"]
scores_embedding = MiniCPM_Embedding_inference(queries, passages)
print(f"MiniCPM-Embedding Scores: {scores_embedding}")
rerank_queries = ["China capital?"]
rerank_passages = [["beijing", "shanghai"]]
scores_reranker = MiniCPM_Reranker_rerank(rerank_queries, rerank_passages)
print(f"MiniCPM-Reranker Scores: {scores_reranker}")
if __name__ == "__main__":
main()

复制代码