Update pages/21_GraphRag.py

pages/21_GraphRag.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import GraphormerForGraphClassification, GraphormerTokenizer
+from transformers import GraphormerForGraphClassification, GraphormerFeatureExtractor
 from datasets import Dataset
 from transformers.models.graphormer.collating_graphormer import preprocess_item, GraphormerDataCollator
 import torch
@@ -14,8 +14,8 @@ def load_model():
         num_classes=2,  # Binary classification (positive/negative sentiment)
         ignore_mismatched_sizes=True,
     )
-    tokenizer = GraphormerTokenizer.from_pretrained("clefourrier/pcqm4mv2_graphormer_base")
-    return model, tokenizer
+    feature_extractor = GraphormerFeatureExtractor.from_pretrained("clefourrier/pcqm4mv2_graphormer_base")
+    return model, feature_extractor
 
 def text_to_graph(text):
     words = text.split()
@@ -36,7 +36,7 @@ def text_to_graph(text):
         "y": [1]  # Placeholder label, will be ignored during inference
     }
 
-def analyze_text(text, model, tokenizer):
+def analyze_text(text, model, feature_extractor):
     graph = text_to_graph(text)
     dataset = Dataset.from_dict({"train": [graph]})
     dataset_processed = dataset.map(preprocess_item, batched=False)
@@ -56,13 +56,13 @@ def analyze_text(text, model, tokenizer):
 
 st.title("Graph-based Text Analysis")
 
-model, tokenizer = load_model()
+model, feature_extractor = load_model()
 
 text_input = st.text_area("Enter text for analysis:", height=200)
 
 if st.button("Analyze Text"):
     if text_input:
-        sentiment, confidence, graph = analyze_text(text_input, model, tokenizer)
+        sentiment, confidence, graph = analyze_text(text_input, model, feature_extractor)
         st.write(f"Sentiment: {sentiment}")
         st.write(f"Confidence: {confidence:.2f}")