Update pages/21_GraphRag.py

pages/21_GraphRag.py

@@ -4,6 +4,22 @@ from transformers import AutoTokenizer, AutoModel
 import torch
 import graphrag
 
+# Diagnostic Section
+st.title("Graphrag Module Investigation")
+
+st.write("Graphrag version:", graphrag.__version__)
+st.write("Contents of graphrag module:")
+st.write(dir(graphrag))
+
+for item in dir(graphrag):
+    st.write(f"Type of {item}: {type(getattr(graphrag, item))}")
+    if callable(getattr(graphrag, item)):
+        st.write(f"Docstring of {item}:")
+        st.write(getattr(graphrag, item).__doc__)
+
+# Main Application Section
+st.title("Graphrag Text Analysis")
+
 @st.cache_resource
 def load_model():
     bert_model_name = "bert-base-uncased"
@@ -11,31 +27,48 @@ def load_model():
     bert_model = AutoModel.from_pretrained(bert_model_name)
 
     # Initialize Graphrag model
-    model = graphrag.GraphRAG(
-        bert_model,
-        num_labels=2,  # Adjust based on your task
-        num_hidden_layers=2,
-        hidden_size=768,
-        intermediate_size=3072,
-    )
+    # Note: This part may need to be adjusted based on the actual structure of graphrag
+    model = None
+    for item in dir(graphrag):
+        if 'model' in item.lower() or 'rag' in item.lower():
+            model_class = getattr(graphrag, item)
+            if callable(model_class):
+                try:
+                    model = model_class(
+                        bert_model,
+                        num_labels=2,  # Adjust based on your task
+                        num_hidden_layers=2,
+                        hidden_size=768,
+                        intermediate_size=3072,
+                    )
+                    break
+                except Exception as e:
+                    st.write(f"Tried initializing {item}, but got error: {str(e)}")
 
-    # If you have a pre-trained Graphrag model, load it here
-    # model.load_state_dict(torch.load('path_to_your_model.pth'))
+    if model is None:
+        st.error("Could not initialize any Graphrag model. Please check the module structure.")
     
     return tokenizer, model
 
 def process_text(text, tokenizer, model):
+    if model is None:
+        return "Model not initialized"
+    
     inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
     with torch.no_grad():
         outputs = model(**inputs)
     # Process outputs based on your specific task
     # This is a placeholder; adjust according to your model's output
-    logits = outputs.logits if hasattr(outputs, 'logits') else outputs
+    if hasattr(outputs, 'logits'):
+        logits = outputs.logits
+    elif isinstance(outputs, torch.Tensor):
+        logits = outputs
+    else:
+        return "Unexpected output format"
+    
     probabilities = torch.softmax(logits, dim=1)
     return probabilities.tolist()[0]
 
-st.title("Graphrag Text Analysis")
-
 tokenizer, model = load_model()
 
 # File uploader