Update pages/21_GraphRag.py

pages/21_GraphRag.py

@@ -1,26 +1,23 @@
-# put code here
-
 import streamlit as st
 import pandas as pd
 from transformers import AutoTokenizer, AutoModel
-from graphrag import GraphragModel, GraphragConfig
 import torch
+from graphrag.models import GraphragForSequenceClassification
+from graphrag.configuration_graphrag import GraphragConfig
 
 @st.cache_resource
 def load_model():
     bert_model_name = "bert-base-uncased"
     tokenizer = AutoTokenizer.from_pretrained(bert_model_name)
-    bert_model = AutoModel.from_pretrained(bert_model_name)
-
-    config = GraphragConfig(
-        bert_model=bert_model,
-        num_labels=2,  # Adjust based on your task
-        num_hidden_layers=2,
-        hidden_size=768,
-        intermediate_size=3072,
-    )
-
-    model = GraphragModel(config)
+    
+    config = GraphragConfig.from_pretrained(bert_model_name)
+    config.num_labels = 2  # Adjust based on your task
+    
+    model = GraphragForSequenceClassification(config)
+    
+    # If you have a pre-trained Graphrag model, load it here
+    # model.load_state_dict(torch.load('path_to_your_model.pth'))
+    
     return tokenizer, model
 
 def process_text(text, tokenizer, model):
@@ -28,7 +25,6 @@ def process_text(text, tokenizer, model):
     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
     probabilities = torch.softmax(logits, dim=1)
     return probabilities.tolist()[0]