Spaces:
Running
Running
| 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" | |
| 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, | |
| max_length=300, | |
| min_length=5, | |
| ) | |
| # 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, | |
| # ) |