Services "corpus_search" and "corpus_count" handle clustering_keys

src_ocaml/grewpy.ml

@@ -2,6 +2,7 @@ open Printf
 open Yojson.Basic.Util
 
 open Conllx
+open Grew_types
 open Libgrew
 
 open Grewpy_utils
@@ -110,8 +111,9 @@ let run_command request =
 
         try
           let conll_corpus = Conllx_corpus.load_list ~quiet:true ~config complete_files in
+          let corpus = Corpus.of_conllx_corpus conll_corpus in
 
-          let index = Global.corpus_add conll_corpus in
+          let index = Global.corpus_add corpus in
           let data = `Assoc [("index", `Int index)] in
           Yojson.Basic.to_string (`Assoc [("status", `String "OK"); ("data", data)])
         with Conllx_error js ->
@@ -123,17 +125,16 @@ let run_command request =
       begin
         try
           let corpus = json |> member "corpus_index" |> to_int |> Global.corpus_get in
-          let position =
+          let graph = 
             match (json |> member "sent_id" |> to_string_option, json |> member "position" |> to_int_option) with
-            | (Some sent_id, _) ->
+            | (Some sent_id, _) -> 
               begin
-                match CCArray.find_idx (fun (id,_) -> id=sent_id) corpus with
-                | Some (i,_) -> i
-                | None -> raise (Error (sprintf "sent_id '%s' not found in corpus" sent_id))
-              end
-            | (_, Some pos) -> pos
+                match Corpus.graph_of_sent_id sent_id corpus with
+                | Some g -> g
+                | None -> raise (Error ("No graph with sent_id: " ^ sent_id))
+              end 
+            | (_, Some pos) -> Corpus.get_graph pos corpus
             | (None, None) -> raise (Error "neither sent_id or pos in the request") in
-          let graph = snd corpus.(position) in
           let data = Graph.to_json graph in
           Yojson.Basic.to_string (`Assoc [("status", `String "OK"); ("data", data)])
         with
@@ -155,8 +156,9 @@ let run_command request =
       begin
         try
           let corpus = json |> member "corpus_index" |> to_int |> Global.corpus_get in
-          let sent_id = Array.to_list (Array.map (fun (id,_) -> `String id) corpus) in
-          Yojson.Basic.to_string (`Assoc [("status", `String "OK"); ("data", `List (sent_id))])
+          let sent_ids = Corpus.fold_left (fun acc sent_id _ -> (`String sent_id) :: acc) [] corpus in
+
+          Yojson.Basic.to_string (`Assoc [("status", `String "OK"); ("data", `List (sent_ids))])
         with
         | Error msg -> json_error msg
       end
@@ -167,24 +169,36 @@ let run_command request =
         try
           let corpus_index = json |> member "corpus_index" |> to_int in
           let request = Request.of_json ~config (json |> member "request") in
+          let clustering_keys = 
+            json 
+            |> member "clustering_keys" 
+            |> to_list 
+            |> List.map (fun x -> Key (to_string x)) in
 
           let corpus = Global.corpus_get corpus_index in
-          let matches = 
-            Array.fold_left
-              (fun acc (id,graph) ->
-                let matching_list = Matching.search_request_in_graph ~config request graph in
-                  (List.map
-                    (fun m -> `Assoc [
-                      ("sent_id", `String id);
-                      ("matching", Matching.to_json request graph m)
-                      ]
-                    ) matching_list
-                  ) @ acc
-              ) [] corpus in
+          let clustered_solutions = 
+            Corpus.search 
+              ~config 
+              [] 
+              (fun sent_id graph matching acc -> `Assoc [
+                ("sent_id", `String sent_id);
+                ("matching", Matching.to_json request graph matching)
+                ] :: acc
+              )
+              request 
+              clustering_keys 
+              corpus in
+          let (json : Yojson.Basic.t) = Clustered.fold_layer
+            (fun x -> `List x)
+            []
+            (fun string_opt sub acc -> (CCOption.get_or ~default:"__undefined__" string_opt, sub) :: acc)
+            (fun x -> `Assoc x)
+            clustered_solutions in
+  
           Yojson.Basic.to_string
             (`Assoc [
                 ("status", `String "OK");
-                ("data", `List matches)
+                ("data", json)
               ])
         with
         | Error msg -> json_error msg
@@ -196,15 +210,23 @@ let run_command request =
         try
           let corpus_index = json |> member "corpus_index" |> to_int in
           let request = Request.of_json ~config (json |> member "request") in
+          let clustering_keys = 
+            json 
+            |> member "clustering_keys" 
+            |> to_list 
+            |> List.map (fun x -> Key (to_string x)) in
 
           let corpus = Global.corpus_get corpus_index in
-          let count = Array.fold_left
-              (fun acc (id,graph) ->
-                 let matching_list = Matching.search_request_in_graph ~config request graph in
-                 (List.length matching_list) + acc
-              ) 0 corpus in
+          let clustered_count = Corpus.search ~config 0 (fun _ _ _ acc -> acc + 1) request clustering_keys corpus in
+          let (json : Yojson.Basic.t) = Clustered.fold_layer
+            (fun x -> `Int x)
+            []
+            (fun string_opt sub acc -> (CCOption.get_or ~default:"__undefined__" string_opt, sub) :: acc)
+            (fun x -> `Assoc x)
+            clustered_count in
+
           Yojson.Basic.to_string
-            (`Assoc [("status", `String "OK"); ("data", `Int count)])
+            (`Assoc [("status", `String "OK"); ("data", json)])
         with
         | Error msg -> json_error msg
       end

src_ocaml/grewpy_utils.ml

@@ -53,24 +53,22 @@ module Global = struct
     incr grs_max;
     grs_map := Int_map.add !grs_max grs !grs_map;
     !grs_max
-  let grs_get index = Int_map.find index !grs_map
+  let grs_get index =
+    try Int_map.find index !grs_map
+    with Not_found -> raise (Error "Reference to an undefined grs")
 
-  (* the [corpus_map] stores corpus loaded by Python *)
-  type corpus = (string * Graph.t) array
-  let (corpus_map: corpus Int_map.t ref) = ref Int_map.empty
+
+  (* the [corpora_map] stores corpora loaded by Python *)
+  let (corpora_map: Corpus.t Int_map.t ref) = ref Int_map.empty
   let corpus_max = ref 0
 
-  let corpus_add conll_corpus =
+  let corpus_add corpus =
     incr corpus_max;
-    let corpus =
-      Array.map
-        (fun (id,conll) -> (id, conll |> Conllx.to_json |> Graph.of_json))
-        (Conllx_corpus.get_data conll_corpus) in
-    corpus_map := Int_map.add !corpus_max corpus !corpus_map;
+    corpora_map := Int_map.add !corpus_max corpus !corpora_map;
     !corpus_max
 
   let corpus_get index =
-    try Int_map.find index !corpus_map 
+    try Int_map.find index !corpora_map 
     with Not_found -> raise (Error "Reference to an undefined corpus")
 end