Update tweet_analyzer.py

tweet_analyzer.py

@@ -2,29 +2,29 @@ import os
 from PyPDF2 import PdfReader
 import pandas as pd
 from dotenv import load_dotenv
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import json
 from datetime import datetime
-from sklearn.decomposition import NMF
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.cluster import KMeans
 import random
-from joblib import Parallel, delayed
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 class TweetDatasetProcessor:
-    def __init__(self):
+    def __init__(self, fine_tuned_model_name):
         load_dotenv()
-        # Load the fine-tuned GPT model and tokenizer
-        self.tokenizer = GPT2Tokenizer.from_pretrained('gpt2')  # Use your fine-tuned model path here
-        self.model = GPT2LMHeadModel.from_pretrained('path_to_finetuned_model')  # Path to your fine-tuned model
         self.tweets = []
-        self.personality_profile = ""
+        self.personality_profile = {}
         self.vectorizer = TfidfVectorizer(stop_words='english')
         self.used_tweets = set()  # Track used tweets to avoid repetition
 
+        # Load fine-tuned model and tokenizer
+        self.model = AutoModelForCausalLM.from_pretrained(fine_tuned_model_name)
+        self.tokenizer = AutoTokenizer.from_pretrained(fine_tuned_model_name)
+
     @staticmethod
     def _process_line(line):
-        """Process a single line."""
+        """Process a single line.""" 
         line = line.strip()
         if not line or line.startswith('http'):  # Skip empty lines and URLs
             return None
@@ -36,7 +36,7 @@ class TweetDatasetProcessor:
         }
 
     def extract_text_from_pdf(self, pdf_path):
-        """Extract text content from PDF file."""
+        """Extract text content from PDF file.""" 
         reader = PdfReader(pdf_path)
         text = ""
         for page in reader.pages:
@@ -44,13 +44,12 @@ class TweetDatasetProcessor:
         return text
 
     def process_pdf_content(self, text):
-        """Process PDF content and clean extracted tweets."""
+        """Process PDF content and clean extracted tweets.""" 
         if not text.strip():
             raise ValueError("The uploaded PDF appears to be empty.")
         
         lines = text.split('\n')
-        # Pass the static method explicitly
-        clean_tweets = Parallel(n_jobs=-1)(delayed(TweetDatasetProcessor._process_line)(line) for line in lines)
+        clean_tweets = [TweetDatasetProcessor._process_line(line) for line in lines]
         self.tweets = [tweet for tweet in clean_tweets if tweet]
 
         if not self.tweets:
@@ -61,16 +60,8 @@ class TweetDatasetProcessor:
         df.to_csv('processed_tweets.csv', index=False)
         return df
 
-    def _extract_mentions(self, text):
-        """Extract mentioned users from tweet."""
-        return [word for word in text.split() if word.startswith('@')]
-
-    def _extract_hashtags(self, text):
-        """Extract hashtags from tweet."""
-        return [word for word in text.split() if word.startswith('#')]
-
     def categorize_tweets(self):
-        """Cluster tweets into categories using KMeans."""
+        """Cluster tweets into categories using KMeans.""" 
         all_tweets = [tweet['content'] for tweet in self.tweets]
         if not all_tweets:
             raise ValueError("No tweets available for clustering.")
@@ -84,7 +75,7 @@ class TweetDatasetProcessor:
         return pd.DataFrame(self.tweets)
 
     def analyze_personality(self, max_tweets=50):
-        """Comprehensive personality analysis using a limited subset of tweets."""
+        """Comprehensive personality analysis using a limited subset of tweets.""" 
         if not self.tweets:
             raise ValueError("No tweets available for personality analysis.")
 
@@ -94,40 +85,16 @@ class TweetDatasetProcessor:
         Tweets for analysis:
         {json.dumps(all_tweets, indent=2)}
         """
-        # Prepare input for the fine-tuned model
-        inputs = self.tokenizer(analysis_prompt, return_tensors="pt", truncation=True, padding=True, max_length=512)
-        
-        try:
-            # Generate response using the fine-tuned model
-            outputs = self.model.generate(inputs['input_ids'], max_length=500)
-            self.personality_profile = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            return self.personality_profile
-        except Exception as e:
-            return f"Error during personality analysis: {str(e)}"
-
-    def analyze_topics(self, n_topics=None):
-        """Extract and identify different topics the author has tweeted about."""
-        all_tweets = [tweet['content'] for tweet in self.tweets]
-        if not all_tweets:
-            return []
-
-        n_topics = n_topics or min(5, len(all_tweets) // 10)
-        tfidf_matrix = self.vectorizer.fit_transform(all_tweets)
-        nmf_model = NMF(n_components=n_topics, random_state=1)
-        nmf_model.fit(tfidf_matrix)
 
-        topics = []
-        for topic_idx, topic in enumerate(nmf_model.components_):
-            topic_words = [self.vectorizer.get_feature_names_out()[i] for i in topic.argsort()[:-n_topics - 1:-1]]
-            topics.append(" ".join(topic_words))
-        return list(set(topics))  # Remove duplicates
+        input_ids = self.tokenizer.encode(analysis_prompt, return_tensors='pt')
+        output = self.model.generate(input_ids, max_length=500, num_return_sequences=1, temperature=0.7)
+        personality_analysis = self.tokenizer.decode(output[0], skip_special_tokens=True)
 
-    def count_tokens(self, text):
-        """Estimate the number of tokens in the given text."""
-        return len(text.split())
+        self.personality_profile = personality_analysis
+        return self.personality_profile
 
     def generate_tweet(self, context="", sample_size=3):
-        """Generate a new tweet by sampling random tweets and avoiding repetition."""
+        """Generate a new tweet by sampling random tweets and avoiding repetition.""" 
         if not self.tweets:
             return "Error: No tweets available for generation."
 
@@ -155,12 +122,9 @@ class TweetDatasetProcessor:
         {', '.join(sampled_contents)}
         **Only generate the tweet. Do not include analysis, explanation, or any other content.**
         """
-        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=512)
+
+        input_ids = self.tokenizer.encode(prompt, return_tensors='pt')
+        output = self.model.generate(input_ids, max_length=150, num_return_sequences=1, temperature=1.0)
+        generated_tweet = self.tokenizer.decode(output[0], skip_special_tokens=True).strip()
         
-        try:
-            # Generate tweet using the fine-tuned model
-            outputs = self.model.generate(inputs['input_ids'], max_length=150)
-            tweet = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            return tweet
-        except Exception as e:
-            return f"Error generating tweet: {str(e)}"
+        return generated_tweet