Update app.py

app.py

@@ -1,3 +1,5 @@
+import warnings
+import numpy as np
 import pandas as pd
 import os
 import json
@@ -7,7 +9,7 @@ import torch
 from sklearn.ensemble import RandomForestClassifier
 from sklearn.model_selection import train_test_split
 from sklearn.preprocessing import OneHotEncoder
-from transformers import AutoModelForSequenceClassification, AutoTokenizer, GPT3LMHeadModel, GPT3Tokenizer, pipeline
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, MegatronLMHeadModel, MegatronTokenizer, pipeline
 from deap import base, creator, tools, algorithms
 import gc
 
@@ -44,11 +46,11 @@ class EmotionalAIAssistant:
         self.emotion_prediction_model = AutoModelForSequenceClassification.from_pretrained("bhadresh-savani/distilbert-base-uncased-emotion")
         self.emotion_prediction_tokenizer = AutoTokenizer.from_pretrained("bhadresh-savani/distilbert-base-uncased-emotion", padding_side='left')
  
-        # Load pre-trained GPT-3 Curie model for text generation
-        self.gpt3_tokenizer = GPT3Tokenizer.from_pretrained('gpt3-curie')
-        self.gpt3_model = GPT3LMHeadModel.from_pretrained('gpt3-curie', device_map='auto')
-      
-        # Enhanced Emotional States
+       # Load pre-trained Megatron-LM model for text generation
+        self.megatron_tokenizer = MegatronTokenizer.from_pretrained('nvidia/megatron-lm-330m')
+        self.megatron_model = MegatronLMHeadModel.from_pretrained('nvidia/megatron-lm-330m', device_map='auto'
+        
+    # Enhanced Emotional States
         self.emotions = {
             'joy': {'percentage': 10, 'motivation': 'positive', 'intensity': 0},
             'pleasure': {'percentage': 10, 'motivation': 'selfish', 'intensity': 0},
@@ -156,6 +158,7 @@ class EmotionalAIAssistant:
         self.emotions['ideal_state']['percentage'] = ideal_state
 
    
+   
     def generate_text(self, prompt, chat_history, emotion=None, max_length=300):
         full_prompt = (
             f"You are Adam, a witty AI assistant who's trying to figure out your own personality and emotions. "
@@ -166,13 +169,13 @@ class EmotionalAIAssistant:
             full_prompt += f"Human: {turn[0]}\nAdam: {turn[1]}\n"
         full_prompt += f"Human: {prompt}\nAdam:"
 
-        input_ids = self.gpt3_tokenizer.encode(full_prompt + self.gpt3_tokenizer.eos_token, return_tensors='pt')
+        input_ids = self.megatron_tokenizer.encode(full_prompt + self.megatron_tokenizer.eos_token, return_tensors='pt')
 
         if torch.cuda.is_available():
             input_ids = input_ids.cuda()
-            self.gpt3_model = self.gpt3_model.cuda()
+            self.megatron_model = self.megatron_model.cuda()
 
-        output = self.gpt3_model.generate(
+        output = self.megatron_model.generate(
             input_ids,
             max_length=len(input_ids[0]) + max_length,
             num_return_sequences=1,
@@ -184,8 +187,8 @@ class EmotionalAIAssistant:
             early_stopping=True,
         )
 
-        generated_text = self.gpt3_tokenizer.decode(output[0], skip_special_tokens=True)
-        return generated_text
+        generated_text = self.megatron_tokenizer.decode(output[0], skip_special_tokens=True)
+        return generated_text 
 
     def predict_emotion(self, context):
         emotion_prediction_pipeline = pipeline('text-classification', model=self.emotion_prediction_model, tokenizer=self.emotion_prediction_tokenizer, top_k=None)