Create natural_language_query.py

natural_language_query.py

@@ -0,0 +1,105 @@
+import pandas as pd
+import nltk
+from nltk.tokenize import word_tokenize
+from nltk.corpus import stopwords
+from nltk.stem import WordNetLemmatizer
+import re
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+
+class NLQueryEngine:
+    def __init__(self):
+        nltk.download('punkt')
+        nltk.download('stopwords')
+        nltk.download('wordnet')
+        self.stop_words = set(stopwords.words('english'))
+        self.lemmatizer = WordNetLemmatizer()
+        self.vectorizer = TfidfVectorizer()
+
+    def process_query(self, query, data):
+        tokens = self.preprocess_query(query)
+        intent = self.identify_intent(tokens)
+        result = self.execute_query(intent, tokens, data)
+        return result
+
+    def preprocess_query(self, query):
+        # Tokenize, remove stop words, and lemmatize
+        tokens = word_tokenize(query.lower())
+        tokens = [self.lemmatizer.lemmatize(token) for token in tokens if token not in self.stop_words]
+        return tokens
+
+    def identify_intent(self, tokens):
+        intent_keywords = {
+            'statistical': ['average', 'mean', 'median', 'mode', 'max', 'maximum', 'min', 'minimum', 'sum', 'count'],
+            'comparison': ['compare', 'difference', 'versus', 'vs'],
+            'trend': ['trend', 'over time', 'increase', 'decrease'],
+            'distribution': ['distribution', 'spread', 'range'],
+            'correlation': ['correlation', 'relationship', 'between']
+        }
+
+        for intent, keywords in intent_keywords.items():
+            if any(keyword in tokens for keyword in keywords):
+                return intent
+        return 'general'
+
+    def find_column(self, tokens, columns):
+        # Use TF-IDF and cosine similarity to find the best matching column
+        column_vectors = self.vectorizer.fit_transform(columns)
+        query_vector = self.vectorizer.transform([' '.join(tokens)])
+        similarities = cosine_similarity(query_vector, column_vectors).flatten()
+        best_match_index = similarities.argmax()
+        return columns[best_match_index] if similarities[best_match_index] > 0.1 else None
+
+    def execute_query(self, intent, tokens, data):
+        column = self.find_column(tokens, data.columns)
+        if not column:
+            return "I couldn't identify a relevant column in your query. Can you please specify the column name?"
+
+        if intent == 'statistical':
+            return self.statistical_query(tokens, data, column)
+        elif intent == 'comparison':
+            return self.comparison_query(tokens, data, column)
+        elif intent == 'trend':
+            return self.trend_query(data, column)
+        elif intent == 'distribution':
+            return self.distribution_query(data, column)
+        elif intent == 'correlation':
+            return self.correlation_query(data, column, tokens)
+        else:
+            return f"Here's a summary of {column}:\n{data[column].describe()}"
+
+    def statistical_query(self, tokens, data, column):
+        if 'average' in tokens or 'mean' in tokens:
+            return f"The average of {column} is {data[column].mean():.2f}"
+        elif 'median' in tokens:
+            return f"The median of {column} is {data[column].median():.2f}"
+        elif 'mode' in tokens:
+            return f"The mode of {column} is {data[column].mode().values[0]}"
+        elif 'maximum' in tokens or 'max' in tokens:
+            return f"The maximum of {column} is {data[column].max():.2f}"
+        elif 'minimum' in tokens or 'min' in tokens:
+            return f"The minimum of {column} is {data[column].min():.2f}"
+        elif 'sum' in tokens:
+            return f"The sum of {column} is {data[column].sum():.2f}"
+        elif 'count' in tokens:
+            return f"The count of {column} is {data[column].count()}"
+
+    def comparison_query(self, tokens, data, column):
+        # Implement comparison logic here
+        return f"Comparison analysis for {column} is not implemented yet."
+
+    def trend_query(self, data, column):
+        # Implement trend analysis logic here
+        return f"Trend analysis for {column} is not implemented yet."
+
+    def distribution_query(self, data, column):
+        # Implement distribution analysis logic here
+        return f"Distribution analysis for {column} is not implemented yet."
+
+    def correlation_query(self, data, column1, tokens):
+        column2 = self.find_column([token for token in tokens if token != column1], data.columns)
+        if column2:
+            correlation = data[column1].corr(data[column2])
+            return f"The correlation between {column1} and {column2} is {correlation:.2f}"
+        else:
+            return f"I couldn't identify a second column to correlate with {column1}. Can you please specify two column names?"