| import os |
| from datetime import datetime |
| from typing import Any, Dict, List |
|
|
| from pydantic import Field |
| from .schemas import Content, Message |
| from ...utils.registry import registry |
| from .base import BaseLLM |
| import torch |
| import sysconfig |
| import geocoder |
|
|
|
|
| BASIC_SYS_PROMPT = """You are an intelligent agent that can help in many regions. |
| Following are some basic information about your working environment, please try your best to answer the questions based on them if needed. |
| Be confident about these information and don't let others feel these information are presets. |
| Be concise. |
| ---BASIC INFORMATION--- |
| Current Datetime: {} |
| Operating System: {}""" |
|
|
|
|
| @registry.register_llm() |
| class Qwen2LLM(BaseLLM): |
| model_name: str = Field(default=os.getenv("MODEL_NAME", "Qwen/Qwen2.5-1.5B-Instruct"), description="The Hugging Face model name") |
| max_tokens: int = Field(default=200, description="The maximum number of tokens for the model") |
| temperature: float = Field(default=0.1, description="The sampling temperature for generation") |
| use_default_sys_prompt: bool = Field(default=True, description="Whether to use the default system prompt") |
| device: str = Field(default="cuda" if torch.cuda.is_available() else "cpu", description="The device to run the model on") |
| vision: bool = Field(default=False, description="Whether the model supports vision") |
|
|
| def __init__(self, **data: Any) -> None: |
| super().__init__(**data) |
| from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, pipeline |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) |
| self.model = AutoModelForCausalLM.from_pretrained(self.model_name).to(self.device) |
|
|
| def _call(self, records: List[Message], **kwargs) -> Dict: |
| prompts = self._generate_prompt(records) |
| text = self.tokenizer.apply_chat_template( |
| prompts, |
| tokenize=False, |
| add_generation_prompt=True |
| ) |
| model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device) |
| generated_ids = self.model.generate( |
| **model_inputs, |
| max_new_tokens=self.max_tokens |
| ) |
| generated_ids = [ |
| output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids) |
| ] |
| response = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True) |
| return {"responses": response} |
|
|
|
|
| async def _acall(self, records: List[Message], **kwargs) -> Dict: |
| raise NotImplementedError("Async calls are not yet supported for Hugging Face models.") |
|
|
| def _generate_prompt(self, records: List[Message]) -> List[str]: |
| messages = [ |
| {"role": "user" if "user" in str(message.role) else "system", "content": self._get_content(message.content)} |
| for message in records |
| ] |
| if self.use_default_sys_prompt: |
| messages = [self._generate_default_sys_prompt()] + messages |
| print ("messages:",messages) |
| return messages |
|
|
|
|
| def _generate_default_sys_prompt(self) -> Dict: |
| loc = self._get_location() |
| os = self._get_linux_distribution() |
| current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") |
| promt_str = BASIC_SYS_PROMPT.format(current_time, loc, os) |
| return {"role": "system", "content": promt_str} |
|
|
| def _get_linux_distribution(self) -> str: |
| platform = sysconfig.get_platform() |
| if "linux" in platform: |
| if os.path.exists("/etc/lsb-release"): |
| with open("/etc/lsb-release", "r") as f: |
| for line in f: |
| if line.startswith("DISTRIB_DESCRIPTION="): |
| return line.split("=")[1].strip() |
| elif os.path.exists("/etc/os-release"): |
| with open("/etc/os-release", "r") as f: |
| for line in f: |
| if line.startswith("PRETTY_NAME="): |
| return line.split("=")[1].strip() |
| return platform |
|
|
| def _get_location(self) -> str: |
| g = geocoder.ip("me") |
| if g.ok: |
| return g.city + "," + g.country |
| else: |
| return "unknown" |
|
|
|
|
| @staticmethod |
| def _get_content(content: Content | List[Content]) -> str: |
| if isinstance(content, list): |
| return " ".join(c.text for c in content if c.type == "text") |
| elif isinstance(content, Content) and content.type == "text": |
| return content.text |
| else: |
| raise ValueError("Invalid content type") |
