1 line
13 KiB
JSON
1 line
13 KiB
JSON
{"repo": "ajjimeno/llm-gp", "files": {"llm.py": {"926a3b726ecb8d510b7db53d0da0a0f11da7499c": {"t": 1785674062.765188, "n": 12231, "h": "16aade81fe1657da", "c": "import importlib\r\nimport os\r\nfrom abc import ABC, abstractmethod\r\nfrom typing import Literal\r\n\r\nimport ollama\r\nimport torch\r\nfrom dotenv import load_dotenv\r\nfrom openai import OpenAI\r\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\r\n\r\nfrom logger_config import getLogger\r\n\r\nload_dotenv()\r\n\r\nlogger = getLogger(__name__)\r\n\r\n\r\ndef get_model(\r\n model_name: Literal[\"qwen\", \"deepseek\", \"ollama\", \"openrouter\"] = \"qwen\"\r\n) -> int:\r\n if model_name == \"qwen\":\r\n return (\r\n Qwen(os.getenv(\"LLM_MODEL_NAME\"))\r\n if os.getenv(\"LLM_MODEL_NAME\", None)\r\n else Qwen()\r\n )\r\n elif model_name == \"deepseek\":\r\n return Openai(\r\n api_key=os.getenv(\"LLM_API_KEY\"), api_url=os.getenv(\"LLM_API_URL\")\r\n )\r\n elif model_name == \"ollama\":\r\n return Ollama()\r\n elif model_name == \"openrouter\":\r\n return OpenRouter()\r\n\r\n raise ValueError(f\"Model {model_name} not available\")\r\n\r\n\r\nclass LLMModel(ABC):\r\n\r\n def __init__(self):\r\n self.model_name = None\r\n self.model = None\r\n self.tokenizer = None\r\n self.bnb_config = None\r\n self.attn = None\r\n self.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\r\n self.do_sampling = False # use the model default as the default here\r\n self.onnx_path = Path(\"./onnx_model\")\r\n\r\n def check_flash_att_compatibility(self) -> bool:\r\n \"\"\"\r\n Check if flash_attention can be used, otherwise default to sdpa\r\n Args:\r\n Returns:\r\n bool value indicating if flash_attn can be used\r\n \"\"\"\r\n\r\n if not torch.cuda.is_available():\r\n print(\"No GPU available, defaulting to CPU\")\r\n return False\r\n\r\n if importlib.util.find_spec(\"flash_attn\") is None:\r\n print(\r\n f\"No package flash_attn available for import. Ensure it is installed and try again!\"\r\n )\r\n return False\r\n\r\n gpu_name = torch.cuda.get_device_name()\r\n gpu_idx = torch.cuda.current_device()\r\n # tuple value representing the minor and major capability of the gpu\r\n gpu_capability = torch.cuda.get_device_capability(gpu_idx)\r\n\r\n print(f\"The following GPU is available: \")\r\n print(f\"\\tname: {gpu_name}\")\r\n print(f\"\\tindex: {gpu_idx}\")\r\n print(f\"\\tCapability: {gpu_capability[0]}.{gpu_capability[1]}\")\r\n\r\n if gpu_capability[0] >= 8: # ampere=8\r\n print(\"gpu support flash_attn\")\r\n return True\r\n else:\r\n print(\"gpu does not support flash_attn\")\r\n return False\r\n\r\n def export_to_onnx(self) -> None:\r\n \"\"\"\r\n Export the model to ONNX for faster inference\r\n Args:\r\n onnx_path: where do we store the ONNX exported model\r\n Return:\r\n None\r\n \"\"\"\r\n\r\n os.makedirs(self.onnx_path, exist_ok=True)\r\n\r\n if self.model is not None and self.tokenizer is not None:\r\n # Make sure the model is completely moved to CPU\r\n # if hasattr(self, 'model'):\r\n # # First unload from GPU\r\n # self.model = self.model.cpu()\r\n\r\n # # Force all parameters to CPU\r\n # for param in self.model.parameters():\r\n # param.data = param.data.cpu()\r\n\r\n # # Force all buffers to CPU\r\n # for buffer in self.model.buffers():\r\n # buffer.data = buffer.data.cpu()\r\n\r\n # Clear CUDA cache\r\n torch.cuda.empty_cache()\r\n with tempfile.TemporaryDirectory() as tmp_dir:\r\n\r\n # perform onnx conversion on CPU as this is memory intensive\r\n # self.model = self.model.cpu()\r\n self.model.save_pretrained(tmp_dir)\r\n self.tokenizer.save_pretrained(tmp_dir)\r\n main_export(\r\n model_name_or_path=tmp_dir,\r\n output=self.onnx_path,\r\n task=\"text-generation\",\r\n opset=14,\r\n device=\"cpu\", # Use CPU for export\r\n no_post_process=True, # Skip post-processing to save memory\r\n trust_remote_code=True,\r\n )\r\n else:\r\n raise ValueError(\r\n \"Ensure there is a model and a tokenizer loaded, before attempting ONNX conversion\"\r\n )\r\n\r\n def load_model(self, device: str) -> None:\r\n \"\"\"\r\n Load a specific LLM model\r\n\r\n Args:\r\n device: the device onto which the model should be loaded. For onnx it is cpu\r\n \"\"\"\r\n\r\n if self.model_name is None:\r\n raise ValueError(\"No model name specified, please set attribute model_name\")\r\n if self.bnb_config is None:\r\n print(\r\n \"No quantization mechanism is implemented. To change this set attribute bnb_config\"\r\n )\r\n if self.attn is None:\r\n print(\r\n \"attention is not specified defaulting to sdpa. To change this set attribute attn\"\r\n )\r\n\r\n self.model = AutoModelForCausalLM.from_pretrained(\r\n self.model_name,\r\n torch_dtype=torch.float16,\r\n device_map=device,\r\n quantization_config=None if self.bnb_config is None else self.bnb_config,\r\n use_sliding_window=False,\r\n do_sample=self.do_sampling, # but we need to look at accuracy\r\n attn_implementation=\"sdpa\" if self.attn is None else self.attn,\r\n )\r\n print(f\"Model loaded on {device}\")\r\n\r\n @abstractmethod\r\n def __call__(self, system_prompt: str, user_prompt: str) -> str:\r\n \"\"\"\r\n Abstract method to process system and user prompts.\r\n\r\n Args:\r\n system_prompt (str): The system prompt.\r\n user_prompt (str): The user prompt.\r\n\r\n Returns:\r\n str: The processed result.\r\n \"\"\"\r\n pass\r\n\r\n def check_flash_att_compatibility(self) -> bool:\r\n \"\"\"\r\n Check if flash_attention can be used, otherwise default to sdpa\r\n Args:\r\n Returns:\r\n bool value indicating if flash_attn can be used\r\n \"\"\"\r\n\r\n if not torch.cuda.is_available():\r\n logger.info(\"No GPU available, defaulting to CPU\")\r\n return False\r\n\r\n if importlib.util.find_spec(\"flash_attn\") is None:\r\n logger.info(\r\n f\"No package flash_attn available for import. Ensure it is installed and try again!\"\r\n )\r\n return False\r\n\r\n gpu_name = torch.cuda.get_device_name()\r\n gpu_idx = torch.cuda.current_device()\r\n # tuple value representing the minor and major capability of the gpu\r\n gpu_capability = torch.cuda.get_device_capability(gpu_idx)\r\n\r\n logger.info(f\"The following GPU is available: \")\r\n logger.info(f\"\\tname: {gpu_name}\")\r\n logger.info(f\"\\tindex: {gpu_idx}\")\r\n logger.info(f\"\\tCapability: {gpu_capability[0]}.{gpu_capability[1]}\")\r\n\r\n if gpu_capability[0] >= 8: # ampere=8\r\n logger.info(\"gpu support flash_attn\")\r\n return True\r\n else:\r\n logger.info(\"gpu does not support flash_attn\")\r\n return False\r\n\r\n\r\nclass Openai(LLMModel):\r\n def __init__(self, api_key, api_url=\"https://api.deepseek.com\"):\r\n self.client = OpenAI(api_key=api_key, base_url=api_url)\r\n\r\n def __call__(self, system_prompt, user_prompt) -> str:\r\n response = self.client.chat.completions.create(\r\n model=\"deepseek-chat\",\r\n messages=[\r\n {\"role\": \"system\", \"content\": system_prompt},\r\n {\"role\": \"user\", \"content\": user_prompt},\r\n ],\r\n stream=False,\r\n )\r\n\r\n return response.choices[0].message.content\r\n\r\n\r\nclass OpenRouter(LLMModel):\r\n def __init__(self):\r\n # print(os.getenv(\"OPENROUTER_API_KEY\"))\r\n self.client = OpenAI(\r\n base_url=\"https://openrouter.ai/api/v1\",\r\n api_key=os.getenv(\"OPENROUTER_API_KEY\"),\r\n )\r\n\r\n def __call__(self, system_prompt, user_prompt) -> str:\r\n print(system_prompt)\r\n print(user_prompt)\r\n response = self.client.chat.completions.create(\r\n model=\"qwen/qwen-2.5-coder-32b-instruct\",\r\n messages=[\r\n {\"role\": \"system\", \"content\": system_prompt},\r\n {\"role\": \"user\", \"content\": user_prompt},\r\n ],\r\n )\r\n\r\n print(response)\r\n logger.info(response.choices[0].message.content)\r\n\r\n return response.choices[0].message.content\r\n\r\n\r\n# Parameter reference\r\n# https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values\r\nclass Ollama(LLMModel):\r\n def __call__(self, system_prompt, user_prompt) -> str:\r\n logger.info(user_prompt)\r\n response = ollama.chat(\r\n \"qwen2.5-coder:32b-instruct-q4_0\",\r\n messages=[\r\n {\"role\": \"system\", \"content\": system_prompt},\r\n {\"role\": \"user\", \"content\": user_prompt},\r\n ],\r\n options={\r\n \"num_ctx\": 32768,\r\n \"num_predict\": 3500,\r\n \"temperature\": 0.0,\r\n },\r\n )\r\n\r\n return response[\"message\"][\"content\"]\r\n\r\n\r\nclass Qwen(LLMModel):\r\n def __init__(\r\n self,\r\n model_name: Literal[\r\n \"Qwen/Qwen2.5-Coder-1.5B-Instruct\", \"Qwen/Qwen2.5-Coder-7B-Instruct\"\r\n ] = \"Qwen/Qwen2.5-Coder-1.5B-Instruct\",\r\n bit_config: Literal[\"8bit\", \"4bit\", \"none\"] = \"4bit\",\r\n use_onnx: bool = True,\r\n ):\r\n super().__init__()\r\n\r\n self.model_name = model_name\r\n if bit_config == \"8bit\":\r\n self.bnb_config = BitsAndBytesConfig(load_in_8bit=True)\r\n elif bit_config == \"4bit\":\r\n self.bnb_config = BitsAndBytesConfig(\r\n load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16\r\n )\r\n\r\n logger.info(f\"Qwen|{model_name}|{bnb_config}\")\r\n\r\n if self.check_flash_att_compatibility():\r\n logger.info(\"Flash attention will be used as the attention mechanism\")\r\n self.attn = \"flash_attention_2\"\r\n else:\r\n logger.info(\r\n \"Unable to run on GPU using flash attention, will run on CPU using sdpa attention mechanism\"\r\n )\r\n self.attn = \"sdpa\"\r\n\r\n if model_name.endswith(\"AWQ\"):\r\n logger.info(f\"Loading AWQ {model_name}\")\r\n self.model = AutoModelForCausalLM.from_pretrained(\r\n model_name,\r\n torch_dtype=torch.float16,\r\n device_map=\"cuda\",\r\n use_sliding_window=False,\r\n attn_implementation=self.attn,\r\n )\r\n else:\r\n self.model = AutoModelForCausalLM.from_pretrained(\r\n model_name,\r\n torch_dtype=torch.float16,\r\n device_map=\"cuda\",\r\n quantization_config=bnb_config,\r\n use_sliding_window=False,\r\n attn_implementation=self.attn,\r\n )\r\n self.tokenizer = AutoTokenizer.from_pretrained(model_name)\r\n\r\n def __call__(self, system_prompt, user_prompt) -> str:\r\n logger.info(user_prompt)\r\n messages = [\r\n {\r\n \"role\": \"system\",\r\n \"content\": system_prompt,\r\n },\r\n {\"role\": \"user\", \"content\": user_prompt},\r\n ]\r\n\r\n text = self.tokenizer.apply_chat_template(\r\n messages, tokenize=False, add_generation_prompt=True\r\n )\r\n\r\n model_inputs = self.tokenizer([text], return_tensors=\"pt\").to(self.model.device)\r\n\r\n with torch.no_grad():\r\n generated_ids = self.model.generate(**model_inputs, max_new_tokens=1500)\r\n\r\n generated_ids = [\r\n output_ids[len(input_ids) :]\r\n for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)\r\n ]\r\n\r\n response = self.tokenizer.batch_decode(\r\n generated_ids, skip_special_tokens=True\r\n )[0]\r\n\r\n return response\r\n"}}}} |