pycityagent 2.0.0a43__cp39-cp39-macosx_11_0_arm64.whl

pycityagent/environment/utils/map_utils.py

@@ -0,0 +1,157 @@
+import math
+from typing import Literal,  Union
+
+import numpy as np
+
+
+def get_angle(x, y):
+    return math.atan2(y, x) * 180 / math.pi
+
+
+def point_on_line_given_distance(start_node, end_node, distance):
+    """
+    Given two points (start_point and end_point) defining a line, and a distance s to travel along the line,
+    return the coordinates of the point reached after traveling s units along the line, starting from start_point.
+
+    Args:
+        start_point (tuple): tuple of (x, y) representing the starting point on the line.
+        end_point (tuple): tuple of (x, y) representing the ending point on the line.
+        distance (float): Distance to travel along the line, starting from start_point.
+
+    Returns:
+        tuple: tuple of (x, y) representing the new point reached after traveling s units along the line.
+    """
+
+    x1, y1 = start_node["x"], start_node["y"]
+    x2, y2 = end_node["x"], end_node["y"]
+
+    # Calculate the slope m and the y-intercept b of the line
+    if x1 == x2:
+        # Vertical line, distance is only along the y-axis
+        return (x1, y1 + distance if distance >= 0 else y1 - abs(distance))
+    else:
+        m = (y2 - y1) / (x2 - x1)
+        b = y1 - m * x1
+
+        # Calculate the direction vector (dx, dy) along the line
+        dx = (x2 - x1) / math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
+        dy = (y2 - y1) / math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
+
+        # Scale the direction vector by the given distance
+        scaled_dx = dx * distance
+        scaled_dy = dy * distance
+
+        # Calculate the new point's coordinates
+        x = x1 + scaled_dx
+        y = y1 + scaled_dy
+
+        return [x, y]
+
+
+def get_key_index_in_lane(
+    nodes: list[dict[str, float]],
+    distance: float,
+    direction: Union[Literal["front"], Literal["back"]],
+) -> int:
+    if direction == "front":
+        _nodes = [n for n in nodes]
+        _index_offset, _index_factor = 0, 1
+    elif direction == "back":
+        _nodes = [n for n in nodes[::-1]]
+        _index_offset, _index_factor = len(_nodes) - 1, -1
+    else:
+        raise ValueError(f"Invalid direction type {direction}!")
+    _lane_points: list[tuple[float, float, float]] = [
+        (n["x"], n["y"], n.get("z", 0)) for n in _nodes
+    ]
+    _line_lengths: list[float] = [0.0 for _ in range(len(_nodes))]
+    _s = 0.0
+    for i, (cur_p, next_p) in enumerate(zip(_lane_points[:-1], _lane_points[1:])):
+        _s += math.hypot(next_p[0] - cur_p[0], next_p[1] - cur_p[1])
+        _line_lengths[i + 1] = _s
+    s = np.clip(distance, _line_lengths[0], _line_lengths[-1])
+    _key_index = 0
+    for (
+        prev_s,
+        cur_s,
+    ) in zip(_line_lengths[:-1], _line_lengths[1:]):
+        if prev_s <= s < cur_s:
+            break
+        _key_index += 1
+    return _index_offset + _index_factor * _key_index
+
+
+def get_xy_in_lane(
+    nodes: list[dict[str, float]],
+    distance: float,
+    direction: Union[Literal["front"], Literal["back"]],
+) -> tuple[float, float]:
+    if direction == "front":
+        _nodes = [n for n in nodes]
+    elif direction == "back":
+        _nodes = [n for n in nodes[::-1]]
+    else:
+        raise ValueError(f"Invalid direction type {direction}!")
+    _lane_points: list[tuple[float, float, float]] = [
+        (n["x"], n["y"], n.get("z", 0)) for n in _nodes
+    ]
+    _line_lengths: list[float] = [0.0 for _ in range(len(_nodes))]
+    _s = 0.0
+    for i, (cur_p, next_p) in enumerate(zip(_lane_points[:-1], _lane_points[1:])):
+        _s += math.hypot(next_p[0] - cur_p[0], next_p[1] - cur_p[1])
+        _line_lengths[i + 1] = _s
+    s = np.clip(distance, _line_lengths[0], _line_lengths[-1])
+    for prev_s, prev_p, cur_s, cur_p in zip(
+        _line_lengths[:-1],
+        _lane_points[:-1],
+        _line_lengths[1:],
+        _lane_points[1:],
+    ):
+        if prev_s <= s < cur_s:
+            _delta_x, _delta_y, _delta_z = [
+                cur_p[_idx] - prev_p[_idx] for _idx in [0, 1, 2]
+            ]
+            _blend_x, _blend_y, _blend_z = [prev_p[_idx] for _idx in [0, 1, 2]]
+            _ratio = (s - prev_s) / (cur_s - prev_s)
+            return (
+                _blend_x + _ratio * _delta_x,
+                _blend_y + _ratio * _delta_y,
+                _blend_z + _ratio * _delta_z,
+            )[:2]
+    return _lane_points[-1][:2]
+
+
+def get_direction_by_s(
+    nodes: list[dict[str, float]],
+    distance: float,
+    direction: Union[Literal["front"], Literal["back"]],
+) -> float:
+    if direction == "front":
+        _nodes = [n for n in nodes]
+    elif direction == "back":
+        _nodes = [n for n in nodes[::-1]]
+    else:
+        raise ValueError(f"Invalid direction type {direction}!")
+    _lane_points: list[tuple[float, float, float]] = [
+        (n["x"], n["y"], n.get("z", 0)) for n in _nodes
+    ]
+    _line_lengths: list[float] = [0.0 for _ in range(len(_nodes))]
+    _line_directions: list[tuple[float, float]] = []
+    _s = 0.0
+    for i, (cur_p, next_p) in enumerate(zip(_lane_points[:-1], _lane_points[1:])):
+        _s += math.hypot(next_p[0] - cur_p[0], next_p[1] - cur_p[1])
+        _line_lengths[i + 1] = _s
+    for i, (cur_p, next_p) in enumerate(zip(_lane_points[:-1], _lane_points[1:])):
+        _direction = math.atan2(next_p[1] - cur_p[1], next_p[0] - cur_p[0])
+        _pitch = math.atan2(
+            next_p[2] - cur_p[2],
+            math.hypot(next_p[0] - cur_p[0], next_p[1] - cur_p[1]),
+        )
+        _line_directions.append((_direction / math.pi * 180, _pitch / math.pi * 180))
+    s = np.clip(distance, _line_lengths[0], _line_lengths[-1])
+    for prev_s, cur_s, direcs in zip(
+        _line_lengths[:-1], _line_lengths[1:], _line_directions
+    ):
+        if prev_s <= s < cur_s:
+            return direcs[0]
+    return _line_directions[-1][0]

pycityagent/environment/utils/port.py

@@ -0,0 +1,11 @@
+import socket
+from contextlib import closing
+
+__all__ = ["find_free_port"]
+
+
+def find_free_port():
+    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]

pycityagent/environment/utils/protobuf.py

@@ -0,0 +1,41 @@
+from collections.abc import Awaitable
+from typing import Any, TypeVar, Union
+
+from google.protobuf.json_format import MessageToDict
+from google.protobuf.message import Message
+
+__all__ = ["parse", "async_parse"]
+
+T = TypeVar("T", bound=Message)
+
+
+def parse(res: T, dict_return: bool) -> Union[dict[str, Any], T]:
+    """
+    将Protobuf返回值转换为dict或者原始值
+    Convert Protobuf return value to dict or original value
+    """
+    if dict_return:
+        return MessageToDict(
+            res,
+            including_default_value_fields=True,
+            preserving_proto_field_name=True,
+            use_integers_for_enums=True,
+        )
+    else:
+        return res
+
+
+async def async_parse(res: Awaitable[T], dict_return: bool) -> Union[dict[str, Any], T]:
+    """
+    将Protobuf await返回值转换为dict或者原始值
+    Convert Protobuf await return value to dict or original value
+    """
+    if dict_return:
+        return MessageToDict(
+            await res,
+            including_default_value_fields=True,
+            preserving_proto_field_name=True,
+            use_integers_for_enums=True,
+        )
+    else:
+        return await res

pycityagent/llm/__init__.py

@@ -0,0 +1,11 @@
+"""LLM相关模块"""
+
+from .embeddings import SentenceEmbedding, SimpleEmbedding
+from .llm import LLM, LLMConfig
+
+__all__ = [
+    "LLM",
+    "LLMConfig",
+    "SentenceEmbedding",
+    "SimpleEmbedding",
+]

pycityagent/llm/embeddings.py

@@ -0,0 +1,231 @@
+import hashlib
+import json
+import os
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from langchain_core.embeddings import Embeddings
+from transformers import AutoModel, AutoTokenizer
+
+__all__ = [
+    "SentenceEmbedding",
+    "SimpleEmbedding",
+]
+
+
+class SentenceEmbedding(Embeddings):
+    def __init__(
+        self,
+        pretrained_model_name_or_path: Union[str, os.PathLike] = "BAAI/bge-m3",
+        max_seq_len: int = 8192,
+        auto_cuda: bool = False,
+        local_files_only: bool = False,
+        cache_dir: str = "./cache",
+        proxies: Optional[dict] = None,
+    ):
+        os.makedirs(cache_dir, exist_ok=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            proxies=proxies,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+        self.model = AutoModel.from_pretrained(
+            pretrained_model_name_or_path,
+            proxies=proxies,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+        self._cuda = auto_cuda and torch.cuda.is_available()
+
+        if self._cuda:
+            self.model = self.model.cuda()
+
+        self.model.eval()
+        self.max_seq_len = max_seq_len
+
+    def _embed(self, texts: list[str]) -> list[list[float]]:
+        # Tokenize sentences
+        encoded_input = self.tokenizer(
+            texts, padding=True, truncation=True, return_tensors="pt"
+        )
+        # for s2p(short query to long passage) retrieval task, add an instruction to query (not add instruction for passages)
+        # encoded_input = tokenizer([instruction + q for q in queries], padding=True, truncation=True, return_tensors='pt')
+
+        # check length of input
+        # assert seq_len <= 8192
+        assert encoded_input["input_ids"].shape[1] <= self.max_seq_len  # type: ignore
+
+        if self._cuda:
+            encoded_input = {k: v.cuda() for k, v in encoded_input.items()}
+        # Compute token embeddings
+        with torch.no_grad():
+            model_output = self.model(**encoded_input)
+            # Perform pooling. In this case, cls pooling.
+            sentence_embeddings = model_output[0][:, 0]
+        # normalize embeddings
+        sentence_embeddings = torch.nn.functional.normalize(
+            sentence_embeddings, p=2, dim=1
+        )
+        if self._cuda:
+            sentence_embeddings = sentence_embeddings.cpu()
+        return sentence_embeddings.tolist()
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Embed documents."""
+        return self._embed(texts)
+
+    def embed_query(self, text: str) -> list[float]:
+        """Embed query text."""
+        return self._embed([text])[0]
+
+
+class SimpleEmbedding(Embeddings):
+    """简单的基于内存的embedding实现
+
+    使用简单的词袋模型(Bag of Words)和TF-IDF来生成文本的向量表示。
+    所有向量都保存在内存中，适用于小规模应用。
+    """
+
+    def __init__(self, vector_dim: int = 128, cache_size: int = 1000):
+        """初始化
+
+        Args:
+            vector_dim: 向量维度
+            cache_size: 缓存大小，超过此大小将清除最早的缓存
+        """
+        self.vector_dim = vector_dim
+        self.cache_size = cache_size
+        self._cache: dict[str, list[float]] = {}
+        self._vocab: dict[str, int] = {}  # 词汇表
+        self._idf: dict[str, float] = {}  # 逆文档频率
+        self._doc_count = 0  # 文档总数
+
+    def _text_to_hash(self, text: str) -> str:
+        """将文本转换为hash值"""
+        return hashlib.md5(text.encode()).hexdigest()
+
+    def _tokenize(self, text: str) -> list[str]:
+        """简单的分词"""
+        # 这里使用简单的空格分词，实际应用中可以使用更复杂的分词方法
+        return text.lower().split()
+
+    def _update_vocab(self, tokens: list[str]):
+        """更新词汇表"""
+        for token in set(tokens):  # 使用set去重
+            if token not in self._vocab:
+                self._vocab[token] = len(self._vocab)
+
+    def _update_idf(self, tokens: list[str]):
+        """更新IDF值"""
+        self._doc_count += 1
+        unique_tokens = set(tokens)
+        for token in unique_tokens:
+            self._idf[token] = self._idf.get(token, 0) + 1
+
+    def _calculate_tf(self, tokens: list[str]) -> dict[str, float]:
+        """计算词频(TF)"""
+        tf = {}
+        total_tokens = len(tokens)
+        for token in tokens:
+            tf[token] = tf.get(token, 0) + 1
+        # 归一化
+        for token in tf:
+            tf[token] /= total_tokens
+        return tf
+
+    def _calculate_tfidf(self, tokens: list[str]) -> list[float]:
+        """计算TF-IDF向量"""
+        vector = np.zeros(self.vector_dim)
+        tf = self._calculate_tf(tokens)
+
+        for token, tf_value in tf.items():
+            if token in self._idf:
+                idf = np.log(self._doc_count / self._idf[token])
+                idx = self._vocab[token] % self.vector_dim  # 使用取模运算来控制向量维度
+                vector[idx] += tf_value * idf
+
+        # L2归一化
+        norm = np.linalg.norm(vector)
+        if norm > 0:
+            vector /= norm
+
+        return list(vector)
+
+    def _embed(self, text: str) -> list[float]:
+        """生成文本的向量表示
+
+        Args:
+            text: 输入文本
+
+        Returns:
+            np.ndarray: 文本的向量表示
+        """
+        # 检查缓存
+        text_hash = self._text_to_hash(text)
+        if text_hash in self._cache:
+            return self._cache[text_hash]
+
+        # 分词
+        tokens = self._tokenize(text)
+        if not tokens:
+            return list(np.zeros(self.vector_dim))
+
+        # 更新词汇表和IDF
+        self._update_vocab(tokens)
+        self._update_idf(tokens)
+
+        # 计算向量
+        vector = self._calculate_tfidf(tokens)
+
+        # 更新缓存
+        if len(self._cache) >= self.cache_size:
+            # 删除最早的缓存
+            oldest_key = next(iter(self._cache))
+            del self._cache[oldest_key]
+        self._cache[text_hash] = vector
+
+        return list(vector)
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Embed documents."""
+        return [self._embed(text) for text in texts]
+
+    def embed_query(self, text: str) -> list[float]:
+        """Embed query text."""
+        return self._embed(text)
+
+    # def save(self, file_path: str):
+    #     """保存模型"""
+    #     state = {
+    #         "vector_dim": self.vector_dim,
+    #         "cache_size": self.cache_size,
+    #         "vocab": self._vocab,
+    #         "idf": self._idf,
+    #         "doc_count": self._doc_count,
+    #     }
+    #     with open(file_path, "w") as f:
+    #         json.dump(state, f)
+
+    # def load(self, file_path: str):
+    #     """加载模型"""
+    #     with open(file_path, "r") as f:
+    #         state = json.load(f)
+    #     self.vector_dim = state["vector_dim"]
+    #     self.cache_size = state["cache_size"]
+    #     self._vocab = state["vocab"]
+    #     self._idf = state["idf"]
+    #     self._doc_count = state["doc_count"]
+    #     self._cache = {}  # 清空缓存
+
+
+if __name__ == "__main__":
+    # se = SentenceEmbedding(
+    #     pretrained_model_name_or_path="ignore/BAAI--bge-m3", cache_dir="ignore"
+    # )
+    se = SimpleEmbedding()
+    print(se.embed_query("hello world"))
+    print(se.embed_query("hello world"))
+    print(se.embed_query("hello world"))
+    print(se.embed_query("hello world"))