Update pages/21_GraphRag.py

pages/21_GraphRag.py

@@ -1,97 +1,58 @@
 import streamlit as st
-import pandas as pd
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
-import graphrag
-import inspect
-
-# Diagnostic Section
-st.title("Graphrag Module Investigation")
-
-st.write("Contents of graphrag module:")
-module_contents = dir(graphrag)
-st.write(module_contents)
-
-st.write("Detailed information about graphrag module contents:")
-for item in module_contents:
-    attr = getattr(graphrag, item)
-    st.write(f"Name: {item}")
-    st.write(f"Type: {type(attr)}")
-    
-    if inspect.isclass(attr):
-        st.write("Class Methods and Attributes:")
-        for name, value in inspect.getmembers(attr):
-            if not name.startswith('_'):  # Exclude private methods/attributes
-                st.write(f"  - {name}: {type(value)}")
-    
-    if callable(attr):
-        st.write("Signature:")
-        st.write(inspect.signature(attr))
-        st.write("Docstring:")
-        st.write(inspect.getdoc(attr))
-    
-    st.write("---")
-
-# Main Application Section
-st.title("Graphrag Text Analysis")
+from collections import Counter
+import nltk
+from nltk.corpus import stopwords
 
 @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)
-
-    # Initialize Graphrag model
-    # Note: This part may need to be adjusted based on the actual structure of graphrag
-    model = None
-    for item in module_contents:
-        if 'model' in item.lower() or 'rag' in item.lower():
-            model_class = getattr(graphrag, item)
-            if inspect.isclass(model_class):
-                try:
-                    # Attempt to initialize the model
-                    # You may need to adjust the parameters based on the actual class signature
-                    model = model_class(bert_model)
-                    st.success(f"Successfully initialized {item}")
-                    break
-                except Exception as e:
-                    st.write(f"Tried initializing {item}, but got error: {str(e)}")
-
-    if model is None:
-        st.error("Could not initialize any Graphrag model. Please check the module structure.")
-    
+    model_name = "distilbert-base-uncased-finetuned-sst-2-english"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
     return tokenizer, model
 
-def process_text(text, tokenizer, model):
-    if model is None:
-        return "Model not initialized"
-    
+def analyze_sentiment(text, tokenizer, model):
     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
-    if hasattr(outputs, 'logits'):
-        logits = outputs.logits
-    elif isinstance(outputs, torch.Tensor):
-        logits = outputs
-    else:
-        return f"Unexpected output format: {type(outputs)}"
+    probabilities = torch.softmax(outputs.logits, dim=1)
+    sentiment = "Positive" if probabilities[0][1] > probabilities[0][0] else "Negative"
+    confidence = probabilities[0][1].item() if sentiment == "Positive" else probabilities[0][0].item()
     
-    probabilities = torch.softmax(logits, dim=1)
-    return probabilities.tolist()[0]
+    return sentiment, confidence
+
+st.title("Text Sentiment Analysis")
 
 tokenizer, model = load_model()
 
-# Text input for single prediction
-text_input = st.text_area("Enter text for analysis:")
+lincoln_text = """Abraham Lincoln (⫽ˈlɪŋkən⫽ LING-kən; February 12, 1809 – April 15, 1865) was an American lawyer, politician, and statesman who served as the 16th president of the United States from 1861 until his assassination in 1865. Lincoln led the United States through the American Civil War, defending the nation as a constitutional union, defeating the insurgent Confederacy, playing a major role in the abolition of slavery, expanding the power of the federal government, and modernizing the U.S. economy.
+Lincoln was born into poverty in a log cabin in Kentucky and was raised on the frontier, mainly in Indiana. He was self-educated and became a lawyer, Whig Party leader, Illinois state legislator, and U.S. representative from Illinois. In 1849, he returned to his successful law practice in Springfield, Illinois. In 1854, angered by the Kansas–Nebraska Act, which opened the territories to slavery, he re-entered politics. He soon became a leader of the new Republican Party. He reached a national audience in the 1858 Senate campaign debates against Stephen A. Douglas. Lincoln ran for president in 1860, sweeping the North to gain victory. Pro-slavery elements in the South viewed his election as a threat to slavery, and Southern states began seceding from the nation. They formed the Confederate States of America, which began seizing federal military bases in the South. A little over one month after Lincoln assumed the presidency, Confederate forces attacked Fort Sumter, a U.S. fort in South Carolina. Following the bombardment, Lincoln mobilized forces to suppress the rebellion and restore the union.
+Lincoln, a moderate Republican, had to navigate a contentious array of factions with friends and opponents from both the Democratic and Republican parties. His allies, the War Democrats and the Radical Republicans, demanded harsh treatment of the Southern Confederates. He managed the factions by exploiting their mutual enmity, carefully distributing political patronage, and by appealing to the American people. Anti-war Democrats (called "Copperheads") despised Lincoln, and some irreconcilable pro-Confederate elements went so far as to plot his assassination. His Gettysburg Address came to be seen as one of the greatest and most influential statements of American national purpose. Lincoln closely supervised the strategy and tactics in the war effort, including the selection of generals, and implemented a naval blockade of the South's trade. He suspended habeas corpus in Maryland and elsewhere, and he averted war with Britain by defusing the Trent Affair. In 1863, he issued the Emancipation Proclamation, which declared the slaves in the states "in rebellion" to be free. It also directed the Army and Navy to "recognize and maintain the freedom of said persons" and to receive them "into the armed service of the United States." Lincoln pressured border states to outlaw slavery, and he promoted the Thirteenth Amendment to the U.S. Constitution, which abolished slavery, except as punishment for a crime.
+Lincoln managed his own successful re-election campaign. He sought to heal the war-torn nation through reconciliation. On April 14, 1865, just five days after the Confederate surrender at Appomattox, he was attending a play at Ford's Theatre in Washington, D.C., with his wife, Mary, when he was fatally shot by Confederate sympathizer John Wilkes Booth. Lincoln is remembered as a martyr and a national hero for his wartime leadership and for his efforts to preserve the Union and abolish slavery. Lincoln is often ranked in both popular and scholarly polls as the greatest president in American history."""
+
+text_input = st.text_area("Enter text for analysis:", value=lincoln_text, height=300)
+
 if st.button("Analyze Text"):
     if text_input:
-        result = process_text(text_input, tokenizer, model)
-        st.write(f"Analysis Result: {result}")
+        sentiment, confidence = analyze_sentiment(text_input, tokenizer, model)
+        st.write(f"Sentiment: {sentiment}")
+        st.write(f"Confidence: {confidence:.2f}")
+        
+        # Additional analysis
+        word_count = len(text_input.split())
+        st.write(f"Word count: {word_count}")
+        
+        # Most common words (excluding stop words)
+        nltk.download('stopwords', quiet=True)
+        stop_words = set(stopwords.words('english'))
+        
+        words = [word.lower() for word in text_input.split() if word.isalnum() and word.lower() not in stop_words]
+        word_freq = Counter(words).most_common(5)
+        
+        st.write("Top 5 most common words (excluding stop words):")
+        for word, freq in word_freq:
+            st.write(f"- {word}: {freq}")
     else:
-        st.write("Please enter some text to analyze.")
-
-# Note about sample data
-st.markdown("Note: To use a CSV file, you would typically upload it and process each row. For simplicity, we're using direct text input in this example.")
+        st.write("Please enter some text to analyze.")