#!/usr/bin/env python3 """ Main application for the Hugging Face Memory Project. Handles conversation interface and memory management. """ import os from dotenv import load_dotenv from transformers import AutoModelForCausalLM, AutoTokenizer import torch from typing import List, Dict, Optional # Import our advanced conversation model try: from src.conversation_model import ConversationModel except ImportError: # Fallback import for direct execution from conversation_model import ConversationModel # Load environment variables (with fallback if dotenv not available) try: load_dotenv() except ImportError: print("⚠️ python-dotenv not available, using default values") class MemoryAI: def __init__(self, use_advanced_model: bool = True): """Initialize the AI model and memory system.""" self.model_name = os.getenv("MODEL_NAME", "gpt2") self.max_memory = int(os.getenv("MAX_MEMORY_ENTRIES", 100)) self.data_dir = os.getenv("DATA_DIR", "data") self.models_dir = os.getenv("MODELS_DIR", "models") # Load generation parameters from environment self.temperature = float(os.getenv("TEMPERATURE", 0.7)) self.max_new_tokens = int(os.getenv("MAX_NEW_TOKENS", 80)) self.top_p = float(os.getenv("TOP_P", 0.9)) self.repetition_penalty = float(os.getenv("REPETITION_PENALTY", 1.2)) # Initialize memory storage self.memories = [] # Initialize conversation model self.use_advanced_model = use_advanced_model self.conversation_model = None if use_advanced_model: try: print("Loading advanced conversation model...") self.conversation_model = ConversationModel() print("✅ Advanced conversation model loaded!") except Exception as e: print(f"❌ Error loading advanced model: {e}") print("Falling back to basic model...") self.use_advanced_model = False # Load basic model as fallback if not self.use_advanced_model: print(f"Loading basic {self.model_name} model...") self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) self.model = AutoModelForCausalLM.from_pretrained(self.model_name) # Move model to GPU if available if torch.cuda.is_available(): self.model = self.model.to('cuda') print("Using CUDA (GPU acceleration)") else: print("Using CPU") print(f"Memory capacity: {self.max_memory} entries") print(f"Generation params - Temp: {self.temperature}, Max tokens: {self.max_new_tokens}") def add_memory(self, memory_text): """Add a memory entry to the system.""" if len(self.memories) >= self.max_memory: self.memories.pop(0) # Remove oldest memory self.memories.append(memory_text) print(f"Memory added. Total memories: {len(self.memories)}") def generate_response(self, prompt, max_new_tokens=80, conversation_history=None): """Generate a response using the AI model with improved quality.""" # Use advanced conversation model if available if self.use_advanced_model and self.conversation_model: try: # Convert memory to conversation history format conv_history = [] if conversation_history: for entry in conversation_history: conv_history.append({"role": entry.get("role", "user"), "content": entry.get("content", entry.get("text", ""))}) # Generate response using advanced model response = self.conversation_model.generate_response(prompt, conv_history) # Add to memories self.add_memory(f"User: {prompt}") self.add_memory(f"AI: {response}") return response except Exception as e: print(f"Advanced model error: {e}") # Fallback to basic model pass # Fallback to basic model # Improved prompt engineering for conversational AI if "microsoft/DialoGPT" in self.model_name: # DialoGPT uses a different format improved_prompt = prompt else: # For other models, use better prompt engineering improved_prompt = f"{prompt}\n\nAssistant:" inputs = self.tokenizer(improved_prompt, return_tensors="pt") # Move inputs to same device as model if hasattr(self, 'model') and next(self.model.parameters()).is_cuda: inputs = {k: v.to('cuda') for k, v in inputs.items()} # Generate with better parameters outputs = self.model.generate( **inputs, max_new_tokens=self.max_new_tokens, temperature=self.temperature, top_p=self.top_p, do_sample=True, # Enable sampling repetition_penalty=self.repetition_penalty, no_repeat_ngram_size=2, # Prevent exact repeats pad_token_id=self.tokenizer.eos_token_id, eos_token_id=self.tokenizer.eos_token_id ) response = self.tokenizer.decode(outputs[0], skip_special_tokens=True) # Extract only the new part of the response response = response[len(improved_prompt):].strip() # Clean up response response = self._clean_response(response) # Fallback for poor responses if not response or len(response.split()) < 2 or response.startswith("I'm"): response = self._generate_fallback_response(prompt) return response def _generate_fallback_response(self, prompt): """Generate a fallback response when the model produces poor output.""" # Simple rule-based responses for common questions prompt_lower = prompt.lower() if "hello" in prompt_lower or "hi" in prompt_lower or "hey" in prompt_lower: return "Hello! I'm MemoryAI, your conversational assistant. How can I help you today?" elif "how are you" in prompt_lower: return "I'm doing well, thank you for asking! As an AI, I'm always ready to chat. How about you?" elif "your name" in prompt_lower: return "I'm MemoryAI! I'm designed to remember our conversations and provide helpful responses." elif "memory" in prompt_lower and ("work" in prompt_lower or "how" in prompt_lower): return "I remember our past conversations and use that context to provide better, more relevant responses. It's like having a conversation with someone who remembers what you've talked about before!" elif "thank" in prompt_lower or "thanks" in prompt_lower: return "You're welcome! I'm happy to help. Is there anything else you'd like to talk about?" elif "joke" in prompt_lower: return "Why don't scientists trust atoms? Because they make up everything!" elif "weather" in prompt_lower: return "I can't check real-time weather, but I hope it's nice where you are! What city are you in?" else: # Generic fallback return "That's an interesting question! As an AI with memory, I can tell you that we've talked about various topics. What would you like to discuss?" def _clean_response(self, response): """Clean up the AI response for better quality.""" # Remove incomplete sentences at the end if response.endswith(('...', '..', '.')) and len(response.split()) < 3: # If it's a very short response ending with dots, keep it pass else: # Remove trailing incomplete words response = response.rstrip('.,;:!?') # Remove excessive repetition words = response.split() if len(words) > 1: # Check for repeated phrases for i in range(min(3, len(words) // 2)): phrase = ' '.join(words[-i-1:-1]) if response.endswith(f" {phrase} {phrase}"): response = response[:-len(f" {phrase}")].rstrip() break # Capitalize first letter if it's a complete sentence if len(response) > 0 and response[0].islower(): response = response[0].upper() + response[1:] # Add punctuation if missing if len(response) > 0 and response[-1] not in ('.', '!', '?'): response += '.' return response def converse(self): """Start a conversation loop with the AI.""" print("🤖 MemoryAI - Advanced Conversation Mode") print("Type 'quit' to exit.") print("Type '!memories' to see recent memories, '!clear' to clear memories") print("Type '!summary' for conversation summary, '!reset' to reset conversation") print("=" * 60) # Initialize conversation history for advanced model conversation_history = [] while True: user_input = input("👤 You: ") if user_input.lower() == 'quit': print("🤖 AI: Goodbye! Have a great day!") break # Handle special commands if user_input.lower() == '!memories': recent_memories = self.get_recent_memories() print("📚 Recent memories:") for i, memory in enumerate(recent_memories, 1): print(f" {i}. {memory}") continue if user_input.lower() == '!clear': self.clear_memories() print("🗑️ Memories cleared!") continue if user_input.lower() == '!summary' and self.use_advanced_model: summary = self.get_conversation_summary() print(f"📊 {summary}") continue if user_input.lower() == '!reset': self.reset_conversation() conversation_history = [] continue if user_input.strip(): # Generate response with conversation history response = self.generate_response(user_input, conversation_history=conversation_history) print(f"🤖 AI: {response}") # Update conversation history conversation_history.append({"role": "user", "content": user_input}) conversation_history.append({"role": "assistant", "content": response}) # Show conversation stats if using advanced model if self.use_advanced_model and self.conversation_model: stats = self.conversation_model.get_conversation_stats() print(f"📊 Topic: {stats['current_topic']} | Emotion: {stats['user_emotion']}") def get_available_models(self): """Get a list of commonly available models.""" models = [ "gpt2", "distilgpt2", "gpt2-medium", "gpt2-large", "EleutherAI/gpt-neo-125M", "facebook/opt-125m", "microsoft/DialoGPT-small", "microsoft/DialoGPT-medium" ] # Add advanced conversation models if self.use_advanced_model: models.extend([ "facebook/blenderbot-400M-distill", "facebook/blenderbot-1B-distill", "microsoft/DialoGPT-large" ]) return models def get_conversation_summary(self) -> str: """Get a summary of the current conversation.""" if not self.use_advanced_model or not self.conversation_model: return "Conversation summary available only with advanced model." return self.conversation_model.get_conversation_summary() def find_similar_memories(self, query: str, top_k: int = 3) -> list: """Find memories similar to the query using semantic search.""" if not self.use_advanced_model or not self.conversation_model: return [] return self.conversation_model.find_similar_conversations(query, top_k) def reset_conversation(self): """Reset the conversation state.""" if self.use_advanced_model and self.conversation_model: self.conversation_model.reset_conversation() print("Conversation reset successfully!") def save_memories(self): """Save memories to a file.""" memory_file = os.path.join(self.data_dir, "memories.txt") with open(memory_file, 'w') as f: for memory in self.memories: f.write(memory + "\n") print(f"Memories saved to {memory_file}") def load_memories(self): """Load memories from a file.""" memory_file = os.path.join(self.data_dir, "memories.txt") if os.path.exists(memory_file): with open(memory_file, 'r') as f: self.memories = [line.strip() for line in f.readlines() if line.strip()] print(f"Loaded {len(self.memories)} memories from {memory_file}") else: print("No existing memories found.") def get_recent_memories(self, count=5): """Get the most recent memories.""" return self.memories[-count:] if self.memories else [] def clear_memories(self): """Clear all memories.""" self.memories = [] print("All memories cleared.") if __name__ == "__main__": ai = MemoryAI() ai.load_memories() # Load existing memories ai.converse() ai.save_memories()