Update app.py

app.py

@@ -9,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 AutoModelForCausalLM, AutoTokenizer, pipeline, AutoModelForSequenceClassification
+from transformers import AutoModelForSequenceClassification, T5ForConditionalGeneration, T5Tokenizer, pipeline
 from deap import base, creator, tools, algorithms
 import gc
 
@@ -44,18 +44,15 @@ emotion_classes = pd.Categorical(df['emotion']).categories
 emotion_prediction_model = AutoModelForSequenceClassification.from_pretrained("bhadresh-savani/distilbert-base-uncased-emotion")
 emotion_prediction_tokenizer = AutoTokenizer.from_pretrained("bhadresh-savani/distilbert-base-uncased-emotion")
 
-# Lazy loading for the fine-tuned language model (DialoGPT)
+# Lazy loading for the fine-tuned language model (FLAN-T5)
 _finetuned_lm_tokenizer = None
 _finetuned_lm_model = None
 
 def get_finetuned_lm_model():
     global _finetuned_lm_tokenizer, _finetuned_lm_model
     if _finetuned_lm_tokenizer is None or _finetuned_lm_model is None:
-        _finetuned_lm_tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
-        _finetuned_lm_tokenizer.padding_side = 'left'  # Set padding to left
-        _finetuned_lm_tokenizer.pad_token = _finetuned_lm_tokenizer.eos_token  # Set pad token to eos token
-        _finetuned_lm_model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium", device_map="auto", low_cpu_mem_usage=True)
-        _finetuned_lm_model.config.pad_token_id = _finetuned_lm_model.config.eos_token_id  # Set pad token id in the model config
+        _finetuned_lm_tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")
+        _finetuned_lm_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base", device_map="auto", low_cpu_mem_usage=True)
     return _finetuned_lm_tokenizer, _finetuned_lm_model
 
 # Enhanced Emotional States
@@ -131,6 +128,7 @@ def evolve_emotions():
     toolbox.register("individual", tools.initCycle, creator.Individual,
                      (toolbox.attr_float,) * (len(emotions) - 1) +
                      (toolbox.attr_intensity,) * len(emotions) +
+                     (lambda: 100,), n(toolbox.attr_intensity,) * len(emotions) +
                      (lambda: 100,), n=1)
     toolbox.register("population", tools.initRepeat, list, toolbox.individual)
     toolbox.register("mate", tools.cxTwoPoint)
@@ -174,42 +172,32 @@ def predict_emotion(context):
 def generate_text(prompt, chat_history, emotion=None, max_length=100):
     finetuned_lm_tokenizer, finetuned_lm_model = get_finetuned_lm_model()
     
-    full_prompt = ""
+    # Prepare the input by concatenating the chat history and the new prompt
+    full_prompt = "You are Adam, an AI assistant. Respond to the following conversation in a natural and engaging way, considering the emotion: " + emotion + "\n\n"
     for turn in chat_history[-5:]:  # Consider last 5 turns for context
-        full_prompt += f"{finetuned_lm_tokenizer.eos_token}{turn[0]}{finetuned_lm_tokenizer.eos_token}{turn[1]}"
-    full_prompt += f"{finetuned_lm_tokenizer.eos_token}{prompt}"
+        full_prompt += f"Human: {turn[0]}\nAdam: {turn[1]}\n"
+    full_prompt += f"Human: {prompt}\nAdam:"
     
-    input_ids = finetuned_lm_tokenizer.encode(full_prompt, return_tensors='pt', truncation=True, max_length=1024)
+    input_ids = finetuned_lm_tokenizer(full_prompt, return_tensors="pt", max_length=512, truncation=True).input_ids
     
     if torch.cuda.is_available():
         input_ids = input_ids.cuda()
         finetuned_lm_model = finetuned_lm_model.cuda()
 
-    # Set up the emotion-specific generation parameters
-    if emotion:
-        # You can adjust these parameters based on the emotion
-        temperature = 0.7
-        top_k = 50
-        top_p = 0.9
-    else:
-        temperature = 1.0
-        top_k = 0
-        top_p = 1.0
-
     # Generate the response
-    output = finetuned_lm_model.generate(
+    outputs = finetuned_lm_model.generate(
         input_ids,
-        max_length=max_length,
+        max_length=max_length + input_ids.shape[1],
         num_return_sequences=1,
         no_repeat_ngram_size=2,
         do_sample=True,
-        temperature=temperature,
-        top_k=top_k,
-        top_p=top_p
+        temperature=0.7,
+        top_k=50,
+        top_p=0.95
     )
 
-    generated_text = finetuned_lm_tokenizer.decode(output[0], skip_special_tokens=True)
-    return generated_text
+    generated_text = finetuned_lm_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_text.strip()
 
 def update_emotion_history(emotion, intensity):
     global emotion_history
@@ -270,8 +258,8 @@ def respond_to_user(user_input, chat_history):
 
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Emotion-Aware AI Chatbot")
-    gr.Markdown("Chat with an AI that understands and responds to emotions.")
+    gr.Markdown("# Adam: Emotion-Aware AI Chatbot")
+    gr.Markdown("Chat with Adam, an AI that understands and responds to emotions.")
     
     chatbot = gr.Chatbot()
     msg = gr.Textbox(label="Type your message here...")