Create graph_utils.py

graph_utils.py

@@ -0,0 +1,94 @@
+from pyvis.network import Network
+import gradio as gr
+from transformers import pipeline
+import os
+
+model_id = models/DReAMy-lib/t5-base-DreamBank-Generation-Act-Char
+
+title = "RE Generation"
+description_GNER = """
+A T5 model tuned to perform text generation, and predict which characters are present in the report. Note that, in the Hall and Van de Castle, the character lists never includes the dreamer. Hence, if you (willingly or not) enter a report that does not contain another character reference, the model will/should (correctly) produce an empty string. Moreover, it is likely that the produced list of CHAR could be longer than the one produced by the SA model, as not all CHAR might be associated with emotions. 
+"""
+
+def get_graph_dict(graph_text):
+    edge_labels = {}
+    if graph_text == "":
+        edge_labels = {("No_Graphs", None):None}
+
+    else:
+        try:
+            for trpl in graph_text[1:-1].split(" | "):
+                h,r,t = trpl[1:-1].split(" # ") 
+                edge_labels[(h,t)] = r
+        except:
+            edge_labels = {("Error", None):None}
+    return edge_labels
+
+def text_to_graph(text):
+    
+    # Use a pipeline as a high-level helper
+    pipe = pipeline(
+        "text2text-generation", 
+        model=model_id, 
+        token=usr_tkn,
+        max_length=300,
+        min_length=10,
+    )
+
+    # generate text graph
+    graph_text = pipe(text)
+
+    graph_text = graph_text[0]["generated_text"]
+
+    # get the nodes: label dict
+    edge_labels = get_graph_dict(graph_text)
+    
+    # create network
+    net = Network(directed=True)
+
+    # nodes & edges
+    for (h, t), r  in edge_labels.items():
+        if (h == "Error") or (h == "No_Graphs"):
+            net.add_node(h, shape="circle")
+            continue
+        else:
+            net.add_node(h, shape="circle")
+            net.add_node(t, shape="circle")
+            net.add_edge(h, t, title=r, label=r)
+
+    # set structure
+    net.repulsion(
+        node_distance=200,
+        central_gravity=0.2,
+        spring_length=200,
+        spring_strength=0.05,
+        damping=0.09
+    )
+
+    net.set_edge_smooth('dynamic') 
+
+    # get html
+    html = net.generate_html()
+    html = html.replace("'", "\"")
+    
+    html_s = f"""<iframe style="width: 100%; height: 600px;margin:0 auto" name="result" allow="midi; geolocation; microphone; camera; 
+    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
+    allow-scripts allow-same-origin allow-popups 
+    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
+    allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""
+
+    return html_s, graph_text
+
+# demo = gr.Interface(
+#         text_to_graph,
+#         inputs=gr.Textbox(label="Text", placeholder="Enter a text here."),
+#         outputs=[gr.HTML(label="Extracted graph"),gr.Textbox(label="Extracted text")]
+#         examples= [
+#         ["I was skating on the outdoor ice pond that used to be across the street from my house. I was not alone, but I did not recognize any of the other people who were skating around. I went through my whole repertoire of jumps, spires, and steps-some of which I can do and some of which I'm not yet sure of. They were all executed flawlessly-some I repeated, some I did only once. I seemed to know that if I went into competition, I would be sure of coming in third because there were only three contestants. Up to that time I hadn't considered it because I hadn't thought I was good enough, but now since everything was going so well, I decided to enter."],
+#         ["I was talking on the telephone to the father of an old friend of mine (boy, 21 years old). We were discussing the party the Saturday night before to which I had invited his son as a guest. I asked him if his son had a good time at the party. He told me not to tell his son that he had told me, but that he had had a good time, except he was a little surprised that I had acted the way I did."], 
+#         ["I was walking alone with my dog in a forest."]
+#         ],
+#         title=title,
+#         description=description_GNER,
+#         cache_examples=True,
+# )