Spaces:
Running
Running
| import os | |
| import torch | |
| import requests | |
| import json | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from configs import ip, api_port, model_path | |
| os.environ['CUDA_LAUNCH_BLOCKING'] = '1' | |
| class Linly: | |
| def __init__(self, mode='api', model_path="Linly-AI/Chinese-LLaMA-2-7B-hf", prefix_prompt = '''请用少于25个字回答以下问题\n\n'''): | |
| # mode = api need | |
| # 定义设置的api的服务器,首先记得运行Linly-api-fast.py 填入ip地址和端口号 | |
| self.url = f"http://{ip}:{api_port}" # local server: http://ip:port | |
| self.headers = { | |
| "Content-Type": "application/json" | |
| } | |
| self.data = { | |
| "question": "北京有什么好玩的地方？" | |
| } | |
| # 全局设定的prompt | |
| self.prefix_prompt = prefix_prompt | |
| self.mode = mode | |
| if mode != 'api': | |
| self.model, self.tokenizer = self.init_model(model_path) | |
| self.history = [] | |
| def init_model(self, path = "Linly-AI/Chinese-LLaMA-2-7B-hf"): | |
| model = AutoModelForCausalLM.from_pretrained(path, device_map="cuda:0", | |
| torch_dtype=torch.bfloat16, trust_remote_code=True) | |
| tokenizer = AutoTokenizer.from_pretrained(path, use_fast=False, trust_remote_code=True) | |
| return model, tokenizer | |
| def generate(self, question, system_prompt=""): | |
| if self.mode != 'api': | |
| self.data["question"] = self.message_to_prompt(question, system_prompt) | |
| inputs = self.tokenizer(self.data["question"], return_tensors="pt").to("cuda:0") | |
| try: | |
| generate_ids = self.model.generate(inputs.input_ids, | |
| max_new_tokens=2048, | |
| do_sample=True, | |
| top_k=20, | |
| top_p=0.84, | |
| temperature=1, | |
| repetition_penalty=1.15, | |
| eos_token_id=2, | |
| bos_token_id=1, | |
| pad_token_id=0) | |
| response = self.tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] | |
| response = response.split("### Response:")[-1] | |
| return response | |
| except: | |
| return "对不起，你的请求出错了，请再次尝试。\nSorry, your request has encountered an error. Please try again.\n" | |
| elif self.mode == 'api': | |
| return self.predict_api(question) | |
| def message_to_prompt(self, message, system_prompt=""): | |
| system_prompt = self.prefix_prompt + system_prompt | |
| for interaction in self.history: | |
| user_prompt, bot_prompt = str(interaction[0]).strip(' '), str(interaction[1]).strip(' ') | |
| system_prompt = f"{system_prompt} User: {user_prompt} Bot: {bot_prompt}" | |
| prompt = f"{system_prompt} ### Instruction:{message.strip()} ### Response:" | |
| return prompt | |
| def predict_api(self, question): | |
| # FastAPI Predict 调用API来进行预测 | |
| self.data["question"] = question | |
| headers = {'Content-Type': 'application/json'} | |
| data = {"prompt": question} | |
| response = requests.post(url=self.url, headers=headers, data=json.dumps(data)) | |
| return response.json()['response'] | |
| def chat(self, system_prompt, message, history): | |
| self.history = history | |
| prompt = self.message_to_prompt(message, system_prompt) | |
| response = self.generate(prompt) | |
| self.history.append([message, response]) | |
| return response, self.history | |
| def clear_history(self): | |
| # 清空历史记录 | |
| self.history = [] | |
| def test(): | |
| llm = Linly(mode='offline',model_path='../Linly-AI/Chinese-LLaMA-2-7B-hf') | |
| answer = llm.generate("如何应对压力？") | |
| print(answer) | |
| if __name__ == '__main__': | |
| test() | |