Spaces:
Sleeping
Sleeping
File size: 4,620 Bytes
1b7e88c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
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") |