Added model list to the UI

app.py

@@ -7,7 +7,7 @@ import rapidjson
 import gradio as gr
 import networkx as nx
 
-from phi3_instruct_graph import Phi3InstructGraph
+from llm_graph import LLMGraph, MODEL_LIST
 from pyvis.network import Network
 from spacy import displacy
 from spacy.tokens import Span
@@ -53,16 +53,16 @@ def handle_text(text=""):
     return " ".join(text.split())
 
 # @spaces.GPU
-def extract_kg(text=""):
+def extract_kg(text="", model=None):
     """
     Extract knowledge graph from text
     """
 
     # Catch empty text
-    if not text:
-        raise gr.Error("⚠️ Text must be provided!")
+    if not text or not model:
+        raise gr.Error("⚠️ Both text and model must be provided!")
     try:
-        model = Phi3InstructGraph()
+        model = LLMGraph(model=model)
         result = model.extract(text)
         return rapidjson.loads(result)
     except Exception as e:
@@ -175,8 +175,8 @@ def create_graph(json_data):
     # Create network visualization
     network = Network(
         width="100%",
-        height="700px",
-        directed=False,
+        # height="700px",
+        height="100vh",
         notebook=False,
         bgcolor="#f8fafc", 
         font_color="#1e293b"
@@ -184,29 +184,29 @@ def create_graph(json_data):
     
     # Configure network display
     network.from_nx(G)
-    network.barnes_hut(
-        gravity=-3000,
-        central_gravity=0.3,
-        spring_length=50,
-        spring_strength=0.001,
-        damping=0.09,
-        overlap=0,
-    )
-    
-    # Customize edge appearance
-    for edge in network.edges:
-        edge['width'] = 2
-        edge['arrows'] = {'to': {'enabled': False, 'type': 'arrow'}}
-        edge['color'] = {'color': '#6366f1', 'highlight': '#4f46e5'}
-        edge['font'] = {'size': 12, 'color': '#4b5563', 'face': 'Arial'}
+    # network.barnes_hut(
+    #     gravity=-3000,
+    #     central_gravity=0.3,
+    #     spring_length=50,
+    #     spring_strength=0.001,
+    #     damping=0.09,
+    #     overlap=0,
+    # )
 
     # Customize node appearance
     for node in network.nodes:
         node['color'] = {'background': '#e0e7ff', 'border': '#6366f1', 'highlight': {'background': '#c7d2fe', 'border': '#4f46e5'}}
         node['font'] = {'size': 14, 'color': '#1e293b'}
         node['shape'] = 'dot'
-        node['size'] = 25
-
+        node['size'] = 20
+    
+    # Customize edge appearance
+    for edge in network.edges:
+        edge['width'] = 4
+        # edge['arrows'] = {'to': {'enabled': False, 'type': 'arrow'}}
+        edge['color'] = {'color': '#6366f1', 'highlight': '#4f46e5'}
+        edge['font'] = {'size': 12, 'color': '#4b5563', 'face': 'Arial'}
+    
     # Generate HTML with iframe to isolate styles
     html = network.generate_html()
     html = html.replace("'", '"')
@@ -217,13 +217,13 @@ def create_graph(json_data):
         allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
         allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""
 
-def process_and_visualize(text, progress=gr.Progress()):
+def process_and_visualize(text, model, progress=gr.Progress()):
     """
     Process text and visualize knowledge graph and entities
     """
 
-    if not text:
-        raise gr.Error("⚠️ Text must be provided!")
+    if not text or not model:
+        raise gr.Error("⚠️ Both text and model must be provided!")
     
     # Check if we're processing the first example for caching
     is_first_example = text == EXAMPLES[0][0]
@@ -242,7 +242,7 @@ def process_and_visualize(text, progress=gr.Progress()):
 
     # Continue with normal processing if cache fails
     progress(0, desc="Starting extraction...")
-    json_data = extract_kg(text)
+    json_data = extract_kg(text, model)
     
     progress(0.5, desc="Creating entity visualization...")
     entities_viz = create_custom_entity_viz(json_data, text)
@@ -275,9 +275,9 @@ def process_and_visualize(text, progress=gr.Progress()):
 EXAMPLES = [
     [handle_text("""The family of Azerbaijan President Ilham Aliyev leads a charmed, glamorous life, thanks in part to financial interests in almost every sector of the economy. 
                  His wife, Mehriban, comes from the privileged and powerful Pashayev family that owns banks, insurance and construction companies, 
-                 a television station and a line of cosmetics. She has led the Heydar Aliyev Foundation, Azerbaijan’s pre-eminent charity behind the construction of schools, 
-                 hospitals and the country’s major sports complex. Their eldest daughter, Leyla, editor of Baku magazine, and her sister, Arzu, 
-                 have financial stakes in a firm that won rights to mine for gold in the western village of Chovdar and Azerfon, the country’s largest mobile phone business. 
+                 a television station and a line of cosmetics. She has led the Heydar Aliyev Foundation, Azerbaijan's pre-eminent charity behind the construction of schools, 
+                 hospitals and the country's major sports complex. Their eldest daughter, Leyla, editor of Baku magazine, and her sister, Arzu, 
+                 have financial stakes in a firm that won rights to mine for gold in the western village of Chovdar and Azerfon, the country's largest mobile phone business. 
                  Arzu is also a significant shareholder in SW Holding, which controls nearly every operation related to Azerbaijan Airlines (“Azal”), from meals to airport taxis. 
                  Both sisters and brother Heydar own property in Dubai valued at roughly $75 million in 2010; 
                  Heydar is the legal owner of nine luxury mansions in Dubai purchased for some $44 million.""")],
@@ -286,16 +286,11 @@ EXAMPLES = [
                  citing lead singer Steven Tyler's unrecoverable vocal cord injury. 
                  The decision comes after months of unsuccessful treatment for Tyler's fractured larynx, which he suffered in September 2023.""")],
     
-    [handle_text("""Chaque matin, les places des villages de Provence s’animent sous le soleil généreux du sud de la France. 
-                 Les marchés provençaux, véritables institutions locales, déploient leurs étals colorés où se mêlent les parfums envoûtants des herbes de Provence, 
-                 des fruits gorgés de soleil et des fromages affinés avec passion. Les producteurs, fiers de leurs récoltes, proposent des tomates juteuses, 
-                 des courgettes encore tièdes du matin, et des olives luisantes baignant dans l’huile d’olive dorée. 
-                 Les discussions vont bon train entre les habitants et les marchands, tandis que les touristes, éblouis par cette abondance, 
-                 goûtent aux spécialités régionales comme la tapenade, les calissons d’Aix ou les délicieuses fougasses à l’huile d’olive. 
-                 Ces marchés ne sont pas seulement des lieux d’achat, mais de véritables scènes de vie où se transmettent les traditions 
-                 culinaires et l’art de vivre à la provençale. Entre deux achats, on s’attarde pour déguster un café serré ou un verre de rosé frais, 
-                 sous le chant des cigales, rappelant que la Provence est avant tout une terre de convivialité et de saveurs authentiques. Une balade dans ces marchés, 
-                 c’est un voyage sensoriel au cœur d’une culture où le temps semble suspendu, entre simplicité et raffinement.""")],
+    [handle_text("""Les jardins du Luxembourg, situés au cœur du sixième arrondissement de Paris, offrent un véritable havre de paix aux citadins pressés. 
+                 Créés au début du dix-septième siècle sur l'initiative de Marie de Médicis, ces jardins à la française s'étendent sur vingt-trois hectares 
+                 et abritent le Palais du Luxembourg, siège du Sénat français. Les promeneurs peuvent y admirer les parterres de fleurs soigneusement entretenus, 
+                 les bassins ornés de statues mythologiques, et les allées bordées de marronniers centenaires. Chaque matin, les jardiniers s'affairent à tailler 
+                 les buis et à arroser les rosiers, perpétuant ainsi une tradition d'excellence horticole qui fait la fierté de la capitale française.""")],
 ]
 
 def generate_first_example_cache():
@@ -308,7 +303,7 @@ def generate_first_example_cache():
 
         try:
             text = EXAMPLES[0][0]
-            model = Phi3InstructGraph()
+            model = MODEL_LIST[0] if MODEL_LIST else None
 
             # Extract data
             json_data = extract_kg(text, model)
@@ -357,16 +352,16 @@ def create_ui():
         gr.Markdown(f"# {TITLE}")
         gr.Markdown(f"{SUBTITLE}")
         
-        # Main content area - redesigned layout
+        # Main content area
         with gr.Row():
             # Left panel - Input controls
             with gr.Column(scale=1):
-                # input_model = gr.Dropdown(
-                #     MODEL_LIST, 
-                #     label="🤖 Select Model",
-                #     info="Choose a model to process your text",
-                #     value=MODEL_LIST[0] if MODEL_LIST else None
-                # )
+                input_model = gr.Dropdown(
+                    MODEL_LIST, 
+                    label="🤖 Select Model",
+                    info="Choose a model to process your text",
+                    value=MODEL_LIST[0] if MODEL_LIST else None
+                )
                 
                 input_text = gr.TextArea(
                     label="📝 Input Text", 
@@ -409,8 +404,7 @@ def create_ui():
         # Functionality
         submit_button.click(
             fn=process_and_visualize, 
-            # inputs=[input_text, input_model],
-            inputs=[input_text],
+            inputs=[input_text, input_model],
             outputs=[output_graph, output_entity_viz, output_json, stats_output]
         )
         
@@ -436,8 +430,8 @@ def create_ui():
         
         # Footer
         gr.Markdown("---")
-        gr.Markdown("📋 **Instructions:** Enter text in any language and click 'Extract & Visualize' to generate a knowledge graph.")
-        gr.Markdown("🛠️ Powered by [Phi-3-mini-128k-instruct-graph](https://huggingface.co/EmergentMethods/Phi-3-mini-128k-instruct-graph)")
+        gr.Markdown("📋 **Instructions:** Enter text in any language, select a model and click `Extract & Visualize` to generate a knowledge graph.")
+        gr.Markdown("🛠️ Powered by [GPT-4.1-mini](https://platform.openai.com/docs/models/gpt-4.1-mini) and [Phi-3-mini-128k-instruct-graph](https://huggingface.co/EmergentMethods/Phi-3-mini-128k-instruct-graph)")
         
     return demo
 

app_old.py

@@ -1,6 +1,6 @@
 # import spaces
 import gradio as gr
-from phi3_instruct_graph import MODEL_LIST, Phi3InstructGraph
+from llm_graph import MODEL_LIST, LLMGraph
 import rapidjson
 from pyvis.network import Network
 import networkx as nx
@@ -43,7 +43,7 @@ def handle_text(text):
 # @spaces.GPU
 def extract(text, model):
     try:
-        model = Phi3InstructGraph(model=model)    
+        model = LLMGraph(model=model)    
         result = model.extract(text)
         return rapidjson.loads(result)
     except Exception as e:

phi3_instruct_graph.py → llm_graph.py

@@ -14,12 +14,20 @@ client = InferenceClient(
     token=api_token
 )
 
-class Phi3InstructGraph:
-    def __init__(self, model = "EmergentMethods/Phi-3-mini-4k-instruct-graph"):
+MODEL_LIST = [
+  "OpenAI/GPT-4.1-mini",
+  "EmergentMethods/Phi-3-mini-128k-instruct-graph",
+]
+
+class LLMGraph:
+    def __init__(self, model="OpenAI/GPT-4.1-mini"):
         """
         Initialize the Phi3InstructGraph with a specified model.
         """
-
+        
+        if model not in MODEL_LIST:
+            raise ValueError(f"Model must be one of {MODEL_LIST}")
+        
         self.model_path = model
 
     def _generate(self, messages):

main.py

@@ -1,6 +1,6 @@
 # import spaces
 import gradio as gr
-from phi3_instruct_graph import MODEL_LIST, Phi3InstructGraph
+from llm_graph import MODEL_LIST, LLMGraph
 import rapidjson
 from pyvis.network import Network
 import networkx as nx
@@ -68,7 +68,7 @@ def handle_text(text):
 # Core extraction function
 # @spaces.GPU
 def extract(text, model):
-    model = Phi3InstructGraph(model=model)    
+    model = LLMGraph(model=model)    
     try:
         result = model.extract(text)
         return rapidjson.loads(result)