update pipeline

.gitignore

@@ -0,0 +1,4 @@
+venv
+flagged
+__pycache__
+.env

app.py

@@ -1,63 +1,113 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
+from gradio_pdf import PDF
+from transformers import pipeline
+
+from configuration import ConfigurationManager
+from kg_retrieval import Neo4JConnection, KnowledgeGraphRetrievalPipeline
+
+pipe = pipeline(
+    "g-retriever-task",
+    model="alfiannajih/g-retriever",
 )
 
+config = ConfigurationManager()
+
+neo4j_config = config.get_neo4j_config()
+neo4j_connection = Neo4JConnection(neo4j_config)
+
+kg_retrieval_config = config.get_kg_retrieval_config()
+kg_retrieval_pipeline = KnowledgeGraphRetrievalPipeline(kg_retrieval_config, neo4j_connection)
+
+def parsing_pdf(pdf_file):
+    pass
+
+def retrieve_kg(pdf_file, description):
+    resume = parsing_pdf(pdf_file)
+    
+    subgraph, textualized_graph = kg_retrieval_pipeline.graph_retrieval_pipeline(resume, description)
+    
+    return subgraph, textualized_graph
+
+def get_feedback(pdf_file, description, max_new_tokens, temperature, top_p):
+    graph, textualized_graph = retrieve_kg(pdf_file, description)
+    
+    inputs = {
+        "inputs": description,
+        "textualized_graph": textualized_graph,
+        "graph": graph
+    }
+
+    generate_kwargs  = {
+        "max_new_tokens": max_new_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "do_sample": True
+    }
+
+    generated = pipe(
+        generate_kwargs=generate_kwargs,
+        **inputs
+    )
+
+    return generated
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        <div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
+            Mentor AI
+        </div>
+        <div style="text-align: center; font-size: 20px; margin-bottom: 20px;">
+            A demo application that provides feedback on your resume using the G-Retriever framework, an LLM powered by a Knowledge Graph.
+        </div>
+        """
+    )
+    with gr.Row(equal_height=True):
+        with gr.Column():
+            pdf_file = PDF(label="Resume Fle")
+        with gr.Column():
+            feedback = gr.Textbox(label="Feedback", interactive=False, lines=15)
+
+    with gr.Row():
+        with gr.Column():
+            description = gr.Textbox(
+                label="Description",
+                lines=3,
+            )
+            submit = gr.Button("Get Feedback")
+
+        with gr.Column():
+            max_new_tokens = gr.Slider(
+                minimum=0,
+                maximum=256,
+                value=128,
+                step=1.0,
+                label="Maximum Output Length",
+                interactive=True
+            )
+            top_p = gr.Slider(
+                minimum=0,
+                maximum=1,
+                value=0.8,
+                step=0.01,
+                label="Top P",
+                interactive=True
+            )
+            temperature = gr.Slider(
+                minimum=0.01,
+                maximum=2,
+                value=1,
+                step=0.01,
+                label="Temperature",
+                interactive=True
+            )
+
+    submit.click(
+        fn=get_feedback,
+        inputs=[pdf_file, description, max_new_tokens, temperature, top_p],
+        outputs=[feedback],
+        show_progress=True
+    )
 
 if __name__ == "__main__":
     demo.launch()

configuration.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+from dotenv import load_dotenv
+import os
+
+load_dotenv()
+
+@dataclass(frozen=True)
+class Neo4jConfig:
+    neo4j_uri: str
+    neo4j_user: str
+    neo4j_password: str
+    neo4j_db: str
+
+@dataclass(frozen=True)
+class KGRetrievalConfig(Neo4jConfig):
+    embedding_model: str
+
+class ConfigurationManager:
+    def __init__(self):
+        pass
+
+    def get_neo4j_config(self) -> Neo4jConfig:
+        config = Neo4jConfig(
+            neo4j_uri=os.getenv("NEO4J_URI"),
+            neo4j_user=os.getenv("NEO4J_USER"),
+            neo4j_password=os.getenv("NEO4J_PASSWORD"),
+            neo4j_db=os.getenv("NEO4J_DB")
+        )
+        return config
+
+    def get_kg_retrieval_config(self) -> KGRetrievalConfig:
+        config = KGRetrievalConfig(
+            neo4j_uri=os.getenv("NEO4J_URI"),
+            neo4j_user=os.getenv("NEO4J_USER"),
+            neo4j_password=os.getenv("NEO4J_PASSWORD"),
+            neo4j_db=os.getenv("NEO4J_DB"),
+            embedding_model="thenlper/gte-base"
+        )
+
+        return config

kg_retrieval.py

@@ -0,0 +1,168 @@
+from neo4j import GraphDatabase
+import pandas as pd
+import torch
+from torch_geometric.data import Data
+
+from configuration import Neo4jConfig, KGRetrievalConfig
+from utils import get_emb_model, retrieval_via_pcst
+
+class Neo4JConnection:
+    def __init__(self, config: Neo4jConfig):
+        self.driver = GraphDatabase.driver(config.neo4j_uri, auth=(config.neo4j_user, config.neo4j_password))
+        self.db = config.neo4j_db
+        self.driver.verify_connectivity()
+
+    def get_session(self):
+        return self.driver.session(database=self.db)
+
+    def get_head_node(self, relation_ids):
+        node_ids = self.driver.execute_query(
+            """
+            MATCH (h)-[r]->()
+            WHERE elementId(r) IN {}
+            RETURN DISTINCT elementId(h) AS id
+            """.format(relation_ids)
+        )
+        nodes = [node.value() for node in node_ids.records]
+        
+        return nodes
+    
+    def get_tail_node(self, relation_ids):
+        node_ids = self.driver.execute_query(
+            """
+            MATCH ()-[r]->(t)
+            WHERE elementId(r) IN {}
+            RETURN DISTINCT elementId(t) AS id
+            """.format(relation_ids)
+        )
+        nodes = [node.value() for node in node_ids.records]
+        
+        return nodes
+    
+    def get_tail_connection_from_head(self, head_ids):
+        relation_ids = self.driver.execute_query(
+            """
+            MATCH (h)-[r]->()
+            WHERE elementId(h) IN {}
+            RETURN DISTINCT elementId(r) AS id LIMIT 50
+            """.format(head_ids)
+        )
+        relations = [relation.value() for relation in relation_ids.records]
+
+        return relations
+
+    def close(self):
+        self.driver.close()
+
+class KnowledgeGraphRetrieval:
+    def __init__(self, config: KGRetrievalConfig, neo4j_connection: Neo4JConnection):
+        self.config = config
+        self.neo4j_connection = neo4j_connection
+        
+        self.embedding_model = get_emb_model(self.config.embedding_model)
+
+    def query_relationship_from_node(self, query, n_query):
+        similar_relations = self.neo4j_connection.driver.execute_query(
+            """
+            CALL db.index.vector.queryNodes('JobTitleIndex', {}, {})
+            YIELD node, score
+            MATCH p=(node)-[r:offered_by]->(connectedNode)
+            RETURN elementId(r) AS id, r.job_description, r.location
+            """.format(n_query, query)
+        )
+        
+        relations = []
+        for relation in similar_relations.records:
+            _id = relation.get("id")
+            text = "Job description: {}".format(relation.get("r.job_description"), relation.get("r.location"))
+
+            relations.append({"rel_id": _id, "text": text})
+        
+        return relations
+    
+class KnowledgeGraphRetrievalPipeline(KnowledgeGraphRetrieval):
+    def __init__(
+        self,
+        config: KGRetrievalConfig,
+        neo4j_connection: Neo4JConnection
+    ):
+        KnowledgeGraphRetrieval.__init__(self, config, neo4j_connection)
+
+    def triples_retrieval(self, resume, desc, top_emb=5):
+        query = resume + [desc]
+        query_emb = self.embedding_model.encode(query, show_progress_bar=False).mean(axis=0).tolist()
+
+        relations = self.query_relationship_from_node(query_emb, top_emb)
+        relation_ids = [r["rel_id"] for r in relations]
+
+        tail_ids = self.neo4j_connection.get_tail_node(relation_ids)
+        tail_connection = self.neo4j_connection.get_tail_connection_from_head(tail_ids)
+        
+        head_ids = self.neo4j_connection.get_head_node(relation_ids)
+        head_connection = self.neo4j_connection.get_tail_connection_from_head(head_ids)
+
+        return relation_ids + tail_connection + head_connection, torch.tensor(query_emb)
+    
+    def build_graph(self, triples, query_emb):
+        with self.neo4j_connection.get_session() as session:
+            result = session.run(
+                """
+                MATCH (h)-[r]->(t)
+                WHERE elementId(r) IN {}
+                RETURN h.name AS h_name, h.embedding AS h_embedding, TYPE(r) AS r_type, r.embedding AS r_embedding, r.description AS job_description, t.embedding AS t_embedding, t.name AS t_name
+                """.format(triples)
+            )
+
+            head_nodes = []
+            tail_nodes = []
+            node_embedding = []
+            node_mapping = {}
+            edge_attr = []
+            edges = []
+            nodes = {}
+
+            for rec in result:
+                if rec.get("h_name") not in node_mapping:
+                    node_embedding.append(rec.get("h_embedding"))
+                    nodes[len(node_mapping)] = rec.get("h_name")
+                    node_mapping[rec.get("h_name")] = len(node_mapping)
+
+                if rec.get("t_name") not in node_mapping:
+                    node_embedding.append(rec.get("t_embedding"))
+                    nodes[len(node_mapping)] = rec.get("t_name")
+                    node_mapping[rec.get("t_name")] = len(node_mapping)
+
+                head_nodes.append(rec.get("h_name"))
+                tail_nodes.append(rec.get("t_name"))
+                edge_attr.append(rec.get("r_embedding"))
+
+                if rec.get("job_description") != None:
+                    textualized_prop = "{}\nJob Description: {}".format(rec.get("r_type"), rec.get("job_description"))
+                else:
+                    textualized_prop = rec.get("r_type")
+
+                edges.append({
+                    "src": node_mapping[rec.get("h_name")],
+                    "edge_attr": textualized_prop,
+                    "dst": node_mapping[rec.get("t_name")]
+                })
+            
+            src = [node_mapping[index] for index in head_nodes]
+            dst = [node_mapping[index] for index in tail_nodes]
+
+            edge_index = torch.tensor([src, dst])
+            edge_attr = torch.tensor(edge_attr)
+
+            graph = Data(x=torch.tensor(node_embedding), edge_index=edge_index, edge_attr=edge_attr)
+            nodes = pd.DataFrame([{'node_id': k, 'node_attr': v} for k, v in nodes.items()], columns=['node_id', 'node_attr'])
+            edges = pd.DataFrame(edges, columns=['src', 'edge_attr', 'dst'])
+            
+            subgraph, desc = retrieval_via_pcst(graph, query_emb, nodes, edges, topk=10, topk_e=3, cost_e=0.5)
+
+            return subgraph, desc
+    
+    def graph_retrieval_pipeline(self, resume, desc, top_emb=5):
+        triples, query_emb = self.triples_retrieval(resume, desc, top_emb)
+        subgraph, textualize_graph = self.build_graph(triples, query_emb)
+        
+        return subgraph, textualize_graph

requirements.txt

@@ -1 +1,12 @@
-huggingface_hub==0.22.2
+transformers==4.44.0
+gradio==4.42.0
+torch_geometric==2.5.3
+torch==2.4.0
+neo4j==5.22.0
+numpy==1.26.4
+pandas==2.2.2
+pcst_fast==1.0.10
+sentence-transformers==3.0.1
+python-dotenv==1.0.1
+sentencepiece==0.2.0
+python-dotenv

utils.py

@@ -0,0 +1,104 @@
+import numpy as np
+from pcst_fast import pcst_fast
+from sentence_transformers import SentenceTransformer
+import torch
+from torch_geometric.data import Data
+
+def retrieval_via_pcst(graph, q_emb, textual_nodes, textual_edges, topk=3, topk_e=3, cost_e=0.5):
+    c = 0.01
+    if len(textual_nodes) == 0 or len(textual_edges) == 0:
+        desc = textual_nodes.to_csv(index=False) + '\n' + textual_edges.to_csv(index=False, columns=['src', 'edge_attr', 'dst'])
+        graph = Data(x=graph.x, edge_index=graph.edge_index, edge_attr=graph.edge_attr, num_nodes=graph.num_nodes)
+        return graph, desc
+
+    root = -1  # unrooted
+    num_clusters = 1
+    pruning = 'gw'
+    verbosity_level = 0
+    if topk > 0:
+        n_prizes = torch.nn.CosineSimilarity(dim=-1)(q_emb, graph.x)
+        topk = min(topk, graph.num_nodes)
+        _, topk_n_indices = torch.topk(n_prizes, topk, largest=True)
+
+        n_prizes = torch.zeros_like(n_prizes)
+        n_prizes[topk_n_indices] = torch.arange(topk, 0, -1).float()
+    else:
+        n_prizes = torch.zeros(graph.num_nodes)
+
+    if topk_e > 0:
+        e_prizes = torch.nn.CosineSimilarity(dim=-1)(q_emb, graph.edge_attr)
+        topk_e = min(topk_e, e_prizes.unique().size(0))
+
+        topk_e_values, _ = torch.topk(e_prizes.unique(), topk_e, largest=True)
+        e_prizes[e_prizes < topk_e_values[-1]] = 0.0
+        last_topk_e_value = topk_e
+        for k in range(topk_e):
+            indices = e_prizes == topk_e_values[k]
+            value = min((topk_e-k)/sum(indices), last_topk_e_value)
+            e_prizes[indices] = value
+            last_topk_e_value = value*(1-c)
+        # reduce the cost of the edges such that at least one edge is selected
+        cost_e = min(cost_e, e_prizes.max().item()*(1-c/2))
+    else:
+        e_prizes = torch.zeros(graph.num_edges)
+
+    costs = []
+    edges = []
+    vritual_n_prizes = []
+    virtual_edges = []
+    virtual_costs = []
+    mapping_n = {}
+    mapping_e = {}
+    for i, (src, dst) in enumerate(graph.edge_index.T.numpy()):
+        prize_e = e_prizes[i]
+        if prize_e <= cost_e:
+            mapping_e[len(edges)] = i
+            edges.append((src, dst))
+            costs.append(cost_e - prize_e)
+        else:
+            virtual_node_id = graph.num_nodes + len(vritual_n_prizes)
+            mapping_n[virtual_node_id] = i
+            virtual_edges.append((src, virtual_node_id))
+            virtual_edges.append((virtual_node_id, dst))
+            virtual_costs.append(0)
+            virtual_costs.append(0)
+            vritual_n_prizes.append(prize_e - cost_e)
+
+    prizes = np.concatenate([n_prizes, np.array(vritual_n_prizes)])
+    num_edges = len(edges)
+    if len(virtual_costs) > 0:
+        costs = np.array(costs+virtual_costs)
+        edges = np.array(edges+virtual_edges)
+
+    vertices, edges = pcst_fast(edges, prizes, costs, root, num_clusters, pruning, verbosity_level)
+
+    selected_nodes = vertices[vertices < graph.num_nodes]
+    selected_edges = [mapping_e[e] for e in edges if e < num_edges]
+    virtual_vertices = vertices[vertices >= graph.num_nodes]
+    if len(virtual_vertices) > 0:
+        virtual_vertices = vertices[vertices >= graph.num_nodes]
+        virtual_edges = [mapping_n[i] for i in virtual_vertices]
+        selected_edges = np.array(selected_edges+virtual_edges)
+
+    edge_index = graph.edge_index[:, selected_edges]
+    selected_nodes = np.unique(np.concatenate([selected_nodes, edge_index[0].numpy(), edge_index[1].numpy()]))
+
+    n = textual_nodes.iloc[selected_nodes]
+    e = textual_edges.iloc[selected_edges]
+    desc = n.to_csv(index=False)+'\n'+e.to_csv(index=False, columns=['src', 'edge_attr', 'dst'])
+
+    mapping = {n: i for i, n in enumerate(selected_nodes.tolist())}
+
+    x = graph.x[selected_nodes]
+    edge_attr = graph.edge_attr[selected_edges]
+    src = [mapping[i] for i in edge_index[0].tolist()]
+    dst = [mapping[i] for i in edge_index[1].tolist()]
+    edge_index = torch.LongTensor([src, dst])
+    data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, num_nodes=len(selected_nodes))
+
+    return data, desc
+
+def get_emb_model(path):
+    model = SentenceTransformer(model_name_or_path=path, trust_remote_code=True)
+
+    return model