Allow maturin to generate docs automatically

Author: torymur

python/outlines_core/__init__.pyi

@@ -21,6 +21,7 @@ macro_rules! type_name {
     };
 }
 
+/// Guide object based on Index.
 #[pyclass(name = "Guide", module = "outlines_core")]
 #[derive(Clone, Debug, PartialEq, Encode, Decode)]
 pub struct PyGuide {
@@ -30,6 +31,7 @@ pub struct PyGuide {
 
 #[pymethods]
 impl PyGuide {
+    /// Creates a Guide object based on Index.
     #[new]
     fn __new__(index: PyIndex) -> Self {
         PyGuide {
@@ -38,10 +40,12 @@ impl PyGuide {
         }
     }
 
+    /// Retrieves current state id of the Guide.
     fn get_state(&self) -> StateId {
         self.state
     }
 
+    /// Gets the list of allowed tokens for the current state.
     fn get_tokens(&self) -> PyResult<Vec<TokenId>> {
         self.index
             .get_allowed_tokens(self.state)
@@ -53,6 +57,7 @@ impl PyGuide {
             )))
     }
 
+    /// Guide moves to the next state provided by the token id and returns a list of allowed tokens.
     fn advance(&mut self, token_id: TokenId) -> PyResult<Vec<TokenId>> {
         match self.index.get_next_state(self.state, token_id) {
             Some(new_state) => {
@@ -66,24 +71,28 @@ impl PyGuide {
         }
     }
 
+    /// Checks if the automaton is in a final state.
     fn is_finished(&self) -> bool {
         self.index.is_final_state(self.state)
     }
 
+    /// Gets the debug string representation of the guide.
     fn __repr__(&self) -> String {
         format!(
             "Guide object with the state={:#?} and {:#?}",
             self.state, self.index
         )
     }
 
+    /// Gets the string representation of the guide.
     fn __str__(&self) -> String {
         format!(
             "Guide object with the state={} and {}",
             self.state, self.index.0
         )
     }
 
+    /// Compares whether two guides are the same.
     fn __eq__(&self, other: &PyGuide) -> bool {
         self == other
     }
@@ -109,12 +118,14 @@ impl PyGuide {
     }
 }
 
+/// Index object based on regex and vocabulary.
 #[pyclass(name = "Index", module = "outlines_core")]
 #[derive(Clone, Debug, PartialEq, Encode, Decode)]
 pub struct PyIndex(Arc<Index>);
 
 #[pymethods]
 impl PyIndex {
+    /// Creates an index from a regex and vocabulary.
     #[new]
     fn __new__(py: Python<'_>, regex: &str, vocabulary: &PyVocabulary) -> PyResult<Self> {
         py.allow_threads(|| {
@@ -124,42 +135,52 @@ impl PyIndex {
         })
     }
 
+    /// Returns allowed tokens in this state.
     fn get_allowed_tokens(&self, state: StateId) -> Option<Vec<TokenId>> {
         self.0.allowed_tokens(&state)
     }
 
+    /// Updates the state.
     fn get_next_state(&self, state: StateId, token_id: TokenId) -> Option<StateId> {
         self.0.next_state(&state, &token_id)
     }
 
+    /// Determines whether the current state is a final state.
     fn is_final_state(&self, state: StateId) -> bool {
         self.0.is_final_state(&state)
     }
 
+    /// Get all final states.
     fn get_final_states(&self) -> HashSet<StateId> {
         self.0.final_states().clone()
     }
 
+    /// Returns the Index as a Python Dict object.
     fn get_transitions(&self) -> HashMap<StateId, HashMap<TokenId, StateId>> {
         self.0.transitions().clone()
     }
 
+    /// Returns the ID of the initial state of the index.
     fn get_initial_state(&self) -> StateId {
         self.0.initial_state()
     }
 
+    /// Gets the debug string representation of the index.
     fn __repr__(&self) -> String {
         format!("{:#?}", self.0)
     }
 
+    /// Gets the string representation of the index.
     fn __str__(&self) -> String {
         format!("{}", self.0)
     }
 
+    /// Compares whether two indexes are the same.
     fn __eq__(&self, other: &PyIndex) -> bool {
         *self.0 == *other.0
     }
 
+    /// Makes a deep copy of the Index.
     fn __deepcopy__(&self, _py: Python<'_>, _memo: Py<PyDict>) -> Self {
         PyIndex(Arc::new((*self.0).clone()))
     }
@@ -185,12 +206,14 @@ impl PyIndex {
     }
 }
 
+/// LLM vocabulary.
 #[pyclass(name = "Vocabulary", module = "outlines_core")]
 #[derive(Clone, Debug, Encode, Decode)]
 pub struct PyVocabulary(Vocabulary);
 
 #[pymethods]
 impl PyVocabulary {
+    /// Creates a vocabulary from eos token id and a map of tokens to token ids.
     #[new]
     fn __new__(py: Python<'_>, eos_token_id: TokenId, map: Py<PyAny>) -> PyResult<PyVocabulary> {
         if let Ok(dict) = map.extract::<HashMap<String, Vec<TokenId>>>(py) {
@@ -213,6 +236,7 @@ impl PyVocabulary {
         }
     }
 
+    /// Creates the vocabulary of a pre-trained model.
     #[staticmethod]
     #[pyo3(signature = (model, revision=None, token=None))]
     fn from_pretrained(
@@ -231,6 +255,7 @@ impl PyVocabulary {
         Ok(PyVocabulary(v))
     }
 
+    /// Inserts new token with token_id or extends list of token_ids if token already present.
     fn insert(&mut self, py: Python<'_>, token: Py<PyAny>, token_id: TokenId) -> PyResult<()> {
         if let Ok(t) = token.extract::<String>(py) {
             return Ok(self.0.try_insert(t, token_id)?);
@@ -244,6 +269,7 @@ impl PyVocabulary {
         )))
     }
 
+    /// Removes a token from vocabulary.
     fn remove(&mut self, py: Python<'_>, token: Py<PyAny>) -> PyResult<()> {
         if let Ok(t) = token.extract::<String>(py) {
             self.0.remove(t);
@@ -259,6 +285,7 @@ impl PyVocabulary {
         )))
     }
 
+    /// Gets token ids of a given token.
     fn get(&self, py: Python<'_>, token: Py<PyAny>) -> PyResult<Option<Vec<TokenId>>> {
         if let Ok(t) = token.extract::<String>(py) {
             return Ok(self.0.token_ids(t.into_bytes()).cloned());
@@ -272,26 +299,32 @@ impl PyVocabulary {
         )))
     }
 
+    /// Gets the end of sentence token id.
     fn get_eos_token_id(&self) -> TokenId {
         self.0.eos_token_id()
     }
 
+    /// Gets the debug string representation of the vocabulary.
     fn __repr__(&self) -> String {
         format!("{:#?}", self.0)
     }
 
+    /// Gets the string representation of the vocabulary.
     fn __str__(&self) -> String {
         format!("{}", self.0)
     }
 
+    /// Compares whether two vocabularies are the same.
     fn __eq__(&self, other: &PyVocabulary) -> bool {
         self.0 == other.0
     }
 
+    /// Returns length of Vocabulary's tokens, excluding EOS token.
     fn __len__(&self) -> usize {
         self.0.tokens().len()
     }
 
+    /// Makes a deep copy of the Vocabulary.
     fn __deepcopy__(&self, _py: Python<'_>, _memo: Py<PyDict>) -> Self {
         PyVocabulary(self.0.clone())
     }
@@ -323,6 +356,7 @@ impl PyVocabulary {
     }
 }
 
+/// Creates regex string from JSON schema with optional whitespace pattern.
 #[pyfunction(name = "build_regex_from_schema")]
 #[pyo3(signature = (json_schema, whitespace_pattern=None))]
 pub fn build_regex_from_schema_py(
@@ -363,6 +397,11 @@ fn register_child_module(parent_module: &Bound<'_, PyModule>) -> PyResult<()> {
     Ok(())
 }
 
+/// This package provides core functionality for structured generation, providing a convenient way to:
+///
+/// - build regular expressions from JSON schemas
+///
+/// - construct an Index object by combining a Vocabulary and regular expression to efficiently map tokens from a given Vocabulary to state transitions in a finite-state automation
 #[pymodule]
 fn outlines_core(m: &Bound<'_, PyModule>) -> PyResult<()> {
     let version = env!("CARGO_PKG_VERSION");