Spaces:
Running
Running
File size: 4,589 Bytes
99e1642 a7aec36 99e1642 a7aec36 99e1642 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
import gradio as gr
from transformers import pipeline
import json
from py2neo import Graph
import time
import requests
import json
import os
# Hugging Face model endpoint for text-generation (GPT-2) and NER (BERT)
llm_endpoint = "https://api-inference.huggingface.co/models/gpt2"
ner_endpoint = "https://api-inference.huggingface.co/models/dbmdz/bert-large-cased-finetuned-conll03-english"
# Define headers for authorization
hf_token = os.getenv('HF_TOKEN')
headers = {
"Authorization": f"Bearer {hf_token}"
}
# Connect to Neo4j database
uri = "neo4j://143.198.173.3:7687" # Replace with your database URI
username = "neo4j" # Replace with your Neo4j username
password = "pickumaterinu" # Replace with your Neo4j password
graph = Graph(uri, auth=(username, password))
print("DB connected")
def query_neo4j(question):
"""Extract entities from a question and retrieve their relationships from Neo4j."""
try:
# Generate entities using NER
ner_response = query_with_retry(ner_endpoint, {"inputs": question})
if 'error' in ner_response:
return f"Error in NER response: {ner_response.get('error')}"
# If NER response is empty or does not contain entities, handle gracefully
if not ner_response:
return "No entities extracted from the question."
# Extract the entities from the NER response
entities = []
for item in ner_response:
if 'word' in item:
entities.append(item['word']) # Add the word to entities list
# Query Neo4j for relationships of extracted entities
result = []
try:
for entity in entities:
query_response = graph.run(
"""
CALL {
WITH $entity AS entity
MATCH (p:Place {id: entity})-[r]->(e)
RETURN p.id AS source_id, type(r) AS relationship, e.id AS target_id
LIMIT 50
}
RETURN source_id, relationship, target_id
""",
entity=entity
)
result.extend([f"{row['source_id']} - {row['relationship']} -> {row['target_id']}" for row in query_response])
except Exception as e:
return f"Error with querying Neo4j: {str(e)}"
return "\n".join(result) if result else None
except Exception as e:
return f"An error occurred while processing the question: {str(e)}"
def query_llm(question):
"""Fetch data from Neo4j and provide an answer using GPT-2."""
graph_data = query_neo4j(question)
if graph_data is not None:
if 'error' in graph_data:
return graph_data
context = f"Graph data:\n{graph_data}"
llm_response = query_with_retry(llm_endpoint, {"inputs": context})
if 'error' in llm_response:
return f"Error in LLM response: {llm_response.get('error')}"
if not llm_response or 'generated_text' not in llm_response[0]:
return "LLM generated an empty response."
return llm_response[0]["generated_text"]
else:
return "No relevant data found"
# query hugging face api with retry for model load
def query_with_retry(endpoint, payload, max_retries=5):
"""Send a request to the Hugging Face API with retry logic in case of model loading delays."""
for _ in range(max_retries):
response = requests.post(endpoint, headers=headers, json=payload)
response_data = response.json()
# Check if the model is ready or if an error occurred
if 'error' in response_data:
if 'currently loading' in response_data['error']:
# Wait for a few seconds before retrying
estimated_time = response_data['error'].split('estimated_time":')[1].split('}')[0]
print(f"Model loading... waiting for {estimated_time} seconds.")
time.sleep(float(estimated_time))
else:
print(f"Error: {response_data['error']}")
break
else:
return response_data
return {"error": "Max retries exceeded, model still loading."}
# Gradio Interface
iface = gr.Interface(
fn=query_llm,
inputs="text",
outputs="text",
live=False,
title="RAG - Neo4j & LLM Integration",
description="Fetches data from Neo4j and generates a response using a Hugging Face GPT-2 model and NER."
)
iface.launch(share=True)
|