From 1f43fffb393a817ebfa4b4276cf418a6886eef63 Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Sun, 17 Nov 2024 14:37:43 -0800
Subject: [PATCH] [README]

---
 README.md                  | 28 ++++++++++++++++++++++++++++
 example.py                 |  2 +-
 pyproject.toml             | 18 +++++++-----------
 srt/__init__.py            |  0
 srt_torch/__init__.py      | 21 +++++++++++++++++++++
 {srt => srt_torch}/main.py |  0
 6 files changed, 57 insertions(+), 12 deletions(-)
 delete mode 100644 srt/__init__.py
 create mode 100644 srt_torch/__init__.py
 rename {srt => srt_torch}/main.py (100%)

diff --git a/README.md b/README.md
index f55ac54..b708a6c 100644
--- a/README.md
+++ b/README.md
@@ -6,3 +6,31 @@
 An open-source non-official community implementation of the model from the paper: Surgical Robot Transformer (SRT): Imitation Learning for Surgical Tasks: https://surgical-robot-transformer.github.io/
 
 
+## Installation
+
+```bash
+pip3 install srt-torch
+```
+
+
+## Usage
+
+```
+
+```
+
+
+## Citation
+
+
+```bibtex
+@misc{kim2024surgicalrobottransformersrt,
+    title={Surgical Robot Transformer (SRT): Imitation Learning for Surgical Tasks}, 
+    author={Ji Woong Kim and Tony Z. Zhao and Samuel Schmidgall and Anton Deguet and Marin Kobilarov and Chelsea Finn and Axel Krieger},
+    year={2024},
+    eprint={2407.12998},
+    archivePrefix={arXiv},
+    primaryClass={cs.RO},
+    url={https://arxiv.org/abs/2407.12998}, 
+}
+```
\ No newline at end of file
diff --git a/example.py b/example.py
index 4aadb9e..4f64344 100644
--- a/example.py
+++ b/example.py
@@ -4,7 +4,7 @@
 
 import torch
 from loguru import logger
-from srt.main import (
+from srt_torch.main import (
     SurgicalRobotTransformer,
     ModelConfig,
     RobotObservation,
diff --git a/pyproject.toml b/pyproject.toml
index 7a8572e..dcecf18 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,15 +3,15 @@ requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
-name = "paper"
+name = "srt-torch"
 version = "0.0.1"
-description = "Paper - Pytorch"
+description = "Surgical Robot Transformer (SRT): Imitation Learning for Surgical Tasks"
 license = "MIT"
 authors = ["Kye Gomez <kye@apac.ai>"]
-homepage = "https://github.com/kyegomez/paper"
-documentation = "https://github.com/kyegomez/paper"  # Add this if you have documentation.
+homepage = "https://github.com/Agora-Lab-AI/SRT"
+documentation = "https://github.com/Agora-Lab-AI/SRT"  # Add this if you have documentation.
 readme = "README.md"  # Assuming you have a README.md
-repository = "https://github.com/kyegomez/paper"
+repository = "https://github.com/Agora-Lab-AI/SRT"
 keywords = ["artificial intelligence", "deep learning", "optimizers", "Prompt Engineering"]
 classifiers = [
     "Development Status :: 4 - Beta",
@@ -23,12 +23,8 @@ classifiers = [
 
 [tool.poetry.dependencies]
 python = "^3.10"
-swarms = "*"
-zetascale = "*"
-
-[tool.poetry.dev-dependencies]
-# Add development dependencies here
-
+loguru = "*"
+torch = "*"
 
 [tool.poetry.group.lint.dependencies]
 ruff = "^0.1.6"
diff --git a/srt/__init__.py b/srt/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/srt_torch/__init__.py b/srt_torch/__init__.py
new file mode 100644
index 0000000..1d168ed
--- /dev/null
+++ b/srt_torch/__init__.py
@@ -0,0 +1,21 @@
+from srt_torch.main import (
+    SurgicalRobotTransformer,
+    ModelConfig,
+    RobotObservation,
+    RobotAction,
+    ImageEncoder,
+    TransformerBlock,
+    ActionPredictor,
+    normalize_rotation_6d,
+)
+
+__all__ = [
+    "SurgicalRobotTransformer",
+    "ModelConfig",
+    "RobotObservation",
+    "RobotAction",
+    "ImageEncoder",
+    "TransformerBlock",
+    "ActionPredictor",
+    "normalize_rotation_6d",
+]
diff --git a/srt/main.py b/srt_torch/main.py
similarity index 100%
rename from srt/main.py
rename to srt_torch/main.py