Adding an initial state for tokensft

lib/SpeechRecognitionFSTs/SpeechRecognitionFSTs.jl

@@ -46,6 +46,8 @@ function tokenfst(
     emission_count = 1
     states = Set(Int[])
     arcs, init, final = [], [], []
+
+    
     
     if isnothing(mapping)
         st = Set(Int[]) 
@@ -56,11 +58,14 @@ function tokenfst(
         end
         mapping = EmissionMapping(length(st), length(tokens))   
     end
+    # println("st "* string(length(st)))
+    # Extra state (infinite loop)
+    init_state = length(mapping) + 1
+    push!(states, init_state)
 
     for (j, token) in enumerate(tokens)
-        offset = length(states)
+        offset = length(states) - 1
         
-        # values = [topo_arc[2] for topo_arc in topo]
         for (i,  topo_arc) in enumerate(topo)
             
             src, dest, weight = offset + topo_arc[1], offset + topo_arc[2], topo_arc[3]
@@ -72,7 +77,32 @@ function tokenfst(
                 dest = dest,
                 weight = S(weight)
             )
+            
+            # Init arc to sources
+            if topo_arc[1] == 1
+                init_arc = Arc(
+                    src = init_state,
+                    isym = 0,
+                    osym = 0,
+                    dest = src,
+                    weight = S(weight)
+                )
+                push!(arcs, init_arc)
+            end
 
+            # Final arc to destinations
+            if topo_arc[1] == length(st)
+                final_arc = Arc(
+                    src = src,
+                    isym = 0,
+                    osym = 0,
+                    dest = init_state,
+                    weight = S(weight)
+                ) 
+                push!(arcs, final_arc)
+            end
+
+            # print("test")
             push!(states, src)
             push!(states, dest)
             push!(arcs, arc)
@@ -89,7 +119,7 @@ function tokenfst(
             push!(states, state)
             push!(final, state => S(weight))
         end
-    end
+        end
     TensorFST(arcs, init, final)
 end