Update app.py

app.py

@@ -1,64 +1,135 @@
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from huggingface_hub import InferenceClient
 import torch
 import gradio as gr
 
 # Load Personality_LM model and tokenizer
-model = AutoModelForSequenceClassification.from_pretrained("KevSun/Personality_LM", ignore_mismatched_sizes=True)
-tokenizer = AutoTokenizer.from_pretrained("KevSun/Personality_LM")
+# The model is pre-trained to evaluate personality traits based on text input
+personality_model = AutoModelForSequenceClassification.from_pretrained("KevSun/Personality_LM", ignore_mismatched_sizes=True)
+personality_tokenizer = AutoTokenizer.from_pretrained("KevSun/Personality_LM")
+
+# Initialize the LLM client (HuggingFaceH4/zephyr-7b-beta)
+llm_client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
 def analyze_personality(text):
-    """Analyze personality traits from input text."""
-    encoded_input = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=512)
-    model.eval()
+    """
+    Analyze personality traits from input text using the Personality_LM model.
+    Args:
+        text (str): The input text used for personality analysis.
+    Returns:
+        dict: A dictionary with personality traits and their corresponding scores.
+    """
+    # Encode the input text for the model
+    encoded_input = personality_tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=512)
+    
+    # Set the model to evaluation mode
+    personality_model.eval()
     with torch.no_grad():
-        outputs = model(**encoded_input)
+        # Perform prediction
+        outputs = personality_model(**encoded_input)
 
+    # Apply softmax to get probabilities for each trait
     predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
     predicted_scores = predictions[0].tolist()
     
+    # Define trait names corresponding to the model's output indices
     trait_names = ["agreeableness", "openness", "conscientiousness", "extraversion", "neuroticism"]
+    
+    # Map traits to their respective scores
     personality_traits = {trait: score for trait, score in zip(trait_names, predicted_scores)}
 
     return personality_traits
 
-def adjust_response(response, traits):
-    """Adjust chatbot response based on personality traits."""
-    if traits["agreeableness"] > 0.5:
-        response = f"{response} 😊 I'm so glad we get along well!"
-    if traits["neuroticism"] > 0.5:
-        response += " But I'm feeling a bit worried about what might happen..."
-    if traits["extraversion"] > 0.5:
-        response += " Let's keep chatting! I love interacting with you."
+def generate_response(user_message, traits):
+    """
+    Generate a chatbot response using the LLM and personality traits.
+    Args:
+        user_message (str): The user's input message.
+        traits (dict): The personality traits with their scores.
+    Returns:
+        str: The chatbot response.
+    """
+    # Create a system message to guide the LLM behavior
+    system_message = (
+        "You are a chatbot with the following personality traits: "
+        f"Agreeableness: {traits['agreeableness']:.2f}, "
+        f"Openness: {traits['openness']:.2f}, "
+        f"Conscientiousness: {traits['conscientiousness']:.2f}, "
+        f"Extraversion: {traits['extraversion']:.2f}, "
+        f"Neuroticism: {traits['neuroticism']:.2f}."
+        " Respond to the user's message in a way that reflects these traits."
+    )
+
+    # Generate a response using the LLM
+    messages = [
+        {"role": "system", "content": system_message},
+        {"role": "user", "content": user_message}
+    ]
+
+    response = ""
+    for message in llm_client.chat_completion(
+        messages,
+        max_tokens=256,
+        stream=True,
+        temperature=0.7,
+        top_p=0.95
+    ):
+        token = message.choices[0].delta.content
+        response += token
+
     return response
 
 def respond(user_message, history, personality_text):
-    """Generate chatbot response based on user input and personality."""
+    """
+    Generate a chatbot response based on user input and personality traits.
+    Args:
+        user_message (str): The user's input message.
+        history (list): A list of message-response tuples to maintain conversation history.
+        personality_text (str): The text defining the chatbot's personality.
+    Returns:
+        tuple: Updated conversation history.
+    """
+    # Analyze personality traits from the provided text
     traits = analyze_personality(personality_text)
-    base_response = f"Hi! You said: {user_message}"
-    final_response = adjust_response(base_response, traits)
+    
+    # Generate a response using the LLM and personality traits
+    final_response = generate_response(user_message, traits)
 
+    # Append the new interaction to the conversation history
     history.append((user_message, final_response))
+    
     return history, history
 
 def personality_demo():
-    """Create the Gradio interface for the chatbot with personality training."""
+    """
+    Create the Gradio interface for the chatbot with personality-based adjustments.
+    Returns:
+        gr.Blocks: The Gradio interface object.
+    """
     with gr.Blocks() as demo:
+        # Header for the chatbot interface
         gr.Markdown("### Personality-Based Chatbot")
 
+        # Textbox for defining personality traits via input text
         personality_textbox = gr.Textbox(
             label="Define Personality Text (Use direct input if no file)", 
             placeholder="Type personality description or paste a sample text here."
         )
 
-        chatbot = gr.Chatbot()
-        msg = gr.Textbox(label="User Input", placeholder="Say something to the chatbot...")
-        clear = gr.Button("Clear Chat")
+        # Chatbot UI elements
+        chatbot = gr.Chatbot()  # Chat display area
+        msg = gr.Textbox(label="User Input", placeholder="Say something to the chatbot...")  # User input box
+        clear = gr.Button("Clear Chat")  # Button to clear chat history
 
+        # Link user input submission to the chatbot response function
         msg.submit(respond, [msg, chatbot, personality_textbox], [chatbot, chatbot])
+        
+        # Link clear button to reset the chat history
         clear.click(lambda: ([], []), None, [chatbot, chatbot])
 
     return demo
 
 if __name__ == "__main__":
+    # Launch the Gradio demo interface
     demo = personality_demo()
     demo.launch()