import os import gradio as gr import networkx as nx import matplotlib.pyplot as plt from langchain_experimental.graph_transformers import LLMGraphTransformer from langchain.chains import GraphQAChain from langchain_core.documents import Document from langchain_community.graphs.networkx_graph import NetworkxEntityGraph from langchain_core.prompts import ChatPromptTemplate from langchain_groq import ChatGroq import pandas as pd from gradio_client import Client import numpy as np from PIL import Image as PILImage import base64 from io import BytesIO # Set the base directory BASE_DIR = os.getcwd() GROQ_API_KEY = os.environ.get('GROQ_API_KEY') # Set up LLM and Flux client llm = ChatGroq(temperature=0, model_name='llama-3.1-8b-instant', groq_api_key=GROQ_API_KEY) flux_client = Client("black-forest-labs/Flux.1-schnell") def create_graph(text): documents = [Document(page_content=text)] llm_transformer_filtered = LLMGraphTransformer(llm=llm) graph_documents_filtered = llm_transformer_filtered.convert_to_graph_documents(documents) graph = NetworkxEntityGraph() for node in graph_documents_filtered[0].nodes: graph.add_node(node.id) for edge in graph_documents_filtered[0].relationships: graph._graph.add_edge( edge.source.id, edge.target.id, relation=edge.type ) return graph, graph_documents_filtered def visualize_graph(graph): plt.figure(figsize=(12, 8)) pos = nx.spring_layout(graph._graph) nx.draw(graph._graph, pos, with_labels=True, node_color='lightblue', node_size=500, font_size=8, font_weight='bold') edge_labels = nx.get_edge_attributes(graph._graph, 'relation') nx.draw_networkx_edge_labels(graph._graph, pos, edge_labels=edge_labels, font_size=6) plt.title("Graph Visualization") plt.axis('off') # Save the plot as an image file graph_viz_path = os.path.join(BASE_DIR, 'graph_visualization.png') plt.savefig(graph_viz_path) plt.close() return graph_viz_path def generate_image(prompt): try: print(f"Generating image with prompt: {prompt}") result = flux_client.predict( prompt=prompt, seed=0, randomize_seed=True, width=1024, height=1024, num_inference_steps=4, api_name="/infer" ) if isinstance(result, tuple) and len(result) > 0 and isinstance(result[0], str): img_str = result[0] # Fix padding if necessary missing_padding = len(img_str) % 4 if missing_padding: img_str += '=' * (4 - missing_padding) img_data = base64.b64decode(img_str) image = PILImage.open(BytesIO(img_data)) elif isinstance(result, tuple) and len(result) > 0 and isinstance(result[0], np.ndarray): image = PILImage.fromarray((result[0] * 255).astype(np.uint8)) elif isinstance(result, PILImage.Image): image = result else: raise ValueError(f"Unexpected result format from flux_client.predict: {type(result)}") image_path = os.path.join(BASE_DIR, 'generated_image.png') image.save(image_path) print(f"Image saved to: {image_path}") return image_path except Exception as e: print(f"Error in generate_image: {str(e)}") import traceback traceback.print_exc() return None def create_relations_table(graph_documents_filtered): df = pd.DataFrame(columns=['node1', 'node2', 'relation']) for edge in graph_documents_filtered[0].relationships: df = pd.concat([df, pd.DataFrame({'node1': [edge.source.id], 'node2': [edge.target.id], 'relation': [edge.type]})], ignore_index=True) return df def process_text(text, question): try: print("Creating graph...") graph, graph_documents_filtered = create_graph(text) print("Setting up GraphQAChain...") graph_rag = GraphQAChain.from_llm( llm=llm, graph=graph, verbose=True ) print("Running question through GraphQAChain...") answer = graph_rag.run(question) print(f"Answer: {answer}") print("Visualizing graph...") graph_viz_path = visualize_graph(graph) print(f"Graph visualization saved to: {graph_viz_path}") print("Creating relations table...") relations_table = create_relations_table(graph_documents_filtered) print("Relations table created") print("Generating summary...") summary_prompt = f"Summarize the following text in one sentence: {text}" summary = llm.invoke(summary_prompt).content print(f"Summary: {summary}") print("Generating image...") image_path = generate_image(summary) if image_path and os.path.exists(image_path): print(f"Generated image saved to: {image_path}") else: print("Failed to generate or save image") return answer, graph_viz_path, relations_table, summary, image_path except Exception as e: print(f"An error occurred in process_text: {str(e)}") import traceback traceback.print_exc() return str(e), None, None, str(e), None def ui_function(text, question): answer, graph_viz_path, relations_table, summary, image_path = process_text(text, question) if isinstance(answer, str) and answer.startswith("An error occurred"): return answer, None, None, answer, None return answer, graph_viz_path, relations_table, summary, image_path # Example text example_text = """The Apollo 11 mission, launched by NASA in July 1969, was the first manned mission to land on the Moon. Commanded by Neil Armstrong and piloted by Buzz Aldrin and Michael Collins, it was the culmination of the Space Race between the United States and the Soviet Union. On July 20, 1969, Armstrong and Aldrin became the first humans to set foot on the lunar surface, while Collins orbited above in the command module.""" # Create Gradio interface with gr.Blocks() as iface: with gr.Row(): with gr.Column(): input_text = gr.Textbox(label="Input Text") question = gr.Textbox(label="Question") example_box = gr.Markdown(f"### Example Paragraph\n\n{example_text}") graph_viz = gr.Image(label="Graph Visualization", type="filepath") with gr.Column(): answer = gr.Textbox(label="Answer") relations_table = gr.Dataframe(label="Relations Table") summary = gr.Textbox(label="Summary") generated_image = gr.Image(label="Generated Image", type="filepath") gr.Button("Run").click(ui_function, inputs=[input_text, question], outputs=[answer, graph_viz, relations_table, summary, generated_image]) footer_text = """ """ gr.Markdown(footer_text) iface.launch()