Upload 2 files

app.py

@@ -120,10 +120,10 @@ def create_graph(json_data):
         height="600px",
         directed=True,
         notebook=False,
-        # bgcolor="#111827", 
-        # font_color="white"
-        bgcolor="#FFFFFF",
-        font_color="#111827"
+        bgcolor="#111827", 
+        font_color="white"
+        # bgcolor="#FFFFFF",
+        # font_color="#111827"
     )
     nt.from_nx(G)
     nt.barnes_hut(
@@ -171,7 +171,6 @@ def process_and_visualize(text, model):
 with gr.Blocks(title="Phi-3 Mini 4k Instruct Graph (by Emergent Methods") as demo:
     gr.Markdown("# Phi-3 Mini 4k Instruct Graph (by Emergent Methods)")
     gr.Markdown("Extract a JSON graph from a text input and visualize it.")
-
     with gr.Row():
         with gr.Column(scale=1):
             input_model = gr.Dropdown(
@@ -206,5 +205,6 @@ with gr.Blocks(title="Phi-3 Mini 4k Instruct Graph (by Emergent Methods") as dem
                 inputs=[input_text, input_model],
                 outputs=[output_graph, output_entity_viz]
             )
+    
 
 demo.launch(share=False)

phi3_instruct_graph.py

@@ -48,10 +48,7 @@ class Phi3InstructGraph:
         return self.pipe(messages, **generation_args)
 
     def _get_messages(self, text):
-        messages = [ 
-            {
-                "role": "system", 
-                "content": dedent("""\n
+        context = dedent("""\n
                     A chat between a curious user and an artificial intelligence Assistant. The Assistant is an expert at identifying entities and relationships in text. The Assistant responds in JSON output only.
 
                     The User provides text in the format:
@@ -79,17 +76,34 @@ class Phi3InstructGraph:
 
                     {"type":"object","properties":{"nodes":{"type":"array","items":{"type":"object","properties":{"id":{"type":"string"},"type":{"type":"string"},"detailed_type":{"type":"string"}},"required":["id","type","detailed_type"],"additionalProperties":false}},"edges":{"type":"array","items":{"type":"object","properties":{"from":{"type":"string"},"to":{"type":"string"},"label":{"type":"string"}},"required":["from","to","label"],"additionalProperties":false}}},"required":["nodes","edges"],"additionalProperties":false}
                         """)
-            }, 
-            {
-                "role": "user", 
-                "content": dedent(f"""\n
-                -------Text begin-------
-                {text}
-                -------Text end-------
-                """)
-            }
-        ]
-        return messages
+        
+        user_message = dedent(f"""\n
+                    -------Text begin-------
+                    {text}
+                    -------Text end-------
+                    """)
+        
+        if self.model == "EmergentMethods/Phi-3-medium-128k-instruct-graph":
+            # model without system message
+            messages = [ 
+                {
+                    "role": "user", 
+                    "content": f"{context}\n Input: {user_message}",
+                }
+            ]
+            return messages
+        else:
+            messages = [ 
+                {
+                    "role": "system", 
+                    "content": context
+                }, 
+                {
+                    "role": "user", 
+                    "content": user_message
+                }
+            ]
+            return messages
     
     
     def extract(self, text):