Update app.py

app.py

@@ -6,16 +6,14 @@ import json
 import random
 import gradio as gr
 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, AutoModelForCausalLM, pipeline
 from deap import base, creator, tools, algorithms
-import gc
 
 warnings.filterwarnings('ignore', category=FutureWarning, module='huggingface_hub.file_download')
 
-# Initialize Example Emotions Dataset
+# Initialize Example Dataset (For Emotion Prediction)
 data = {
     'context': [
         'I am happy', 'I am sad', 'I am angry', 'I am excited', 'I am calm',
@@ -44,18 +42,10 @@ 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-medium)
-_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:
-        model_name = "microsoft/DialoGPT-medium"
-        _finetuned_lm_tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left')
-        _finetuned_lm_model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", low_cpu_mem_usage=True)
-        _finetuned_lm_tokenizer.pad_token = _finetuned_lm_tokenizer.eos_token
-    return _finetuned_lm_tokenizer, _finetuned_lm_model
+# Load pre-trained LLM model and tokenizer for response generation
+response_model_name = "microsoft/DialoGPT-medium"
+response_tokenizer = AutoTokenizer.from_pretrained(response_model_name)
+response_model = AutoModelForCausalLM.from_pretrained(response_model_name)
 
 # Enhanced Emotional States
 emotions = {
@@ -200,8 +190,13 @@ def generate_response(context):
     # Ensure pad_token_id is a tensor
     pad_token_id = torch.tensor(tokenizer.pad_token_id)
 
-    outputs = model.generate(inputs, max_length=500, num_return_sequences=1, pad_token_id=pad_token_id.item())
+    outputs = model.generate(inputs, max_length=500, num_return_sequences=1, pad_token_id=pad_token_id.item(), eos_token_id=tokenizer.eos_token_id)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Ensure the response does not repeat the input
+    if context in response:
+        response = response.replace(context, '').strip()
+    
     return response
 
 def handle_conversation(user_input):
@@ -218,7 +213,7 @@ with gr.Blocks() as demo:
     user_input = gr.Textbox(label="User Input")
     response = gr.Textbox(label="Bot Response")
     submit = gr.Button("Submit")
-    submit.click(update_ui, user_input, response)
+    submit.click(update_ui, inputs=[user_input], outputs=[response])
 
 if __name__ == "__main__":
     demo.launch(share=True)