Update pages/21_GraphRag.py

pages/21_GraphRag.py

@@ -1,21 +1,19 @@
 import streamlit as st
 import pandas as pd
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, BertForSequenceClassification, BertConfig
 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)
     
-    config = GraphragConfig.from_pretrained(bert_model_name)
+    config = BertConfig.from_pretrained(bert_model_name)
     config.num_labels = 2  # Adjust based on your task
     
-    model = GraphragForSequenceClassification(config)
+    model = BertForSequenceClassification.from_pretrained(bert_model_name, config=config)
     
-    # If you have a pre-trained Graphrag model, load it here
+    # If you have a pre-trained model, load it here
     # model.load_state_dict(torch.load('path_to_your_model.pth'))
     
     return tokenizer, model
@@ -29,7 +27,7 @@ def process_text(text, tokenizer, model):
     probabilities = torch.softmax(logits, dim=1)
     return probabilities.tolist()[0]
 
-st.title("Graphrag Text Analysis")
+st.title("BERT Text Analysis")
 
 tokenizer, model = load_model()
 
@@ -59,4 +57,4 @@ if st.button("Analyze Text"):
         st.write("Please enter some text to analyze.")
 
 # Add a link to sample data
-st.markdown("[Download Sample CSV](https://raw.githubusercontent.com/your_username/your_repo/main/sample_data.csv)")
+st.markdown("[Download Sample CSV](https://raw.githubusercontent.com/your_username/your_repo/main/sample_data.csv)")