add in changes to make app work as well as demo

app.py

@@ -4,10 +4,22 @@ import tempfile
 import os
 from utils import SocialGraphManager, SuggestionGenerator
 
-# Initialize the social graph manager and suggestion generator
+# Define available models
+AVAILABLE_MODELS = {
+    "distilgpt2": "DistilGPT2 (Fast, smaller model)",
+    "gpt2": "GPT-2 (Medium size, better quality)",
+    "google/gemma-3-1b-it": "Gemma 3 1B-IT (Small, instruction-tuned)",
+    "Qwen/Qwen1.5-0.5B": "Qwen 1.5 0.5B (Very small, efficient)",
+    "Qwen/Qwen1.5-1.8B": "Qwen 1.5 1.8B (Small, good quality)",
+    "TinyLlama/TinyLlama-1.1B-Chat-v1.0": "TinyLlama 1.1B (Small, chat-tuned)",
+    "microsoft/phi-3-mini-4k-instruct": "Phi-3 Mini (Small, instruction-tuned)",
+    "microsoft/phi-2": "Phi-2 (Small, high quality for size)",
+}
+
+# Initialize the social graph manager
 social_graph = SocialGraphManager("social_graph.json")
 
-# Initialize the suggestion generator with distilgpt2
+# Initialize the suggestion generator with distilgpt2 (default)
 suggestion_generator = SuggestionGenerator("distilgpt2")
 
 # Test the model to make sure it's working
@@ -23,7 +35,8 @@ if not suggestion_generator.model_loaded:
 try:
     whisper_model = whisper.load_model("tiny")
     whisper_loaded = True
-except Exception:
+except Exception as e:
+    print(f"Error loading Whisper model: {e}")
     whisper_loaded = False
 
 
@@ -90,16 +103,55 @@ def on_person_change(person_id):
     return context_info, phrases_text, topics
 
 
-def generate_suggestions(person_id, user_input, suggestion_type, selected_topic=None):
+def change_model(model_name):
+    """Change the language model used for generation.
+
+    Args:
+        model_name: The name of the model to use
+
+    Returns:
+        A status message about the model change
+    """
+    global suggestion_generator
+
+    print(f"Changing model to: {model_name}")
+
+    # Check if we need to change the model
+    if model_name == suggestion_generator.model_name:
+        return f"Already using model: {model_name}"
+
+    # Try to load the new model
+    success = suggestion_generator.load_model(model_name)
+
+    if success:
+        return f"Successfully switched to model: {model_name}"
+    else:
+        return f"Failed to load model: {model_name}. Using fallback responses instead."
+
+
+def generate_suggestions(
+    person_id,
+    user_input,
+    suggestion_type,
+    selected_topic=None,
+    model_name="distilgpt2",
+    temperature=0.7,
+):
     """Generate suggestions based on the selected person and user input."""
     print(
-        f"Generating suggestions with: person_id={person_id}, user_input={user_input}, suggestion_type={suggestion_type}, selected_topic={selected_topic}"
+        f"Generating suggestions with: person_id={person_id}, user_input={user_input}, "
+        f"suggestion_type={suggestion_type}, selected_topic={selected_topic}, "
+        f"model={model_name}, temperature={temperature}"
     )
 
     if not person_id:
         print("No person_id provided")
         return "Please select who you're talking to first."
 
+    # Make sure we're using the right model
+    if model_name != suggestion_generator.model_name:
+        change_model(model_name)
+
     person_context = social_graph.get_person_context(person_id)
     print(f"Person context: {person_context}")
 
@@ -160,7 +212,7 @@ def generate_suggestions(person_id, user_input, suggestion_type, selected_topic=
             print(f"Generating suggestion {i+1}/3")
             try:
                 suggestion = suggestion_generator.generate_suggestion(
-                    person_context, user_input, temperature=0.7
+                    person_context, user_input, temperature=temperature
                 )
                 print(f"Generated suggestion: {suggestion}")
                 suggestions.append(suggestion)
@@ -168,11 +220,13 @@ def generate_suggestions(person_id, user_input, suggestion_type, selected_topic=
                 print(f"Error generating suggestion: {e}")
                 suggestions.append("Error generating suggestion")
 
-        result = "### AI-Generated Responses:\n\n"
+        result = (
+            f"### AI-Generated Responses (using {suggestion_generator.model_name}):\n\n"
+        )
         for i, suggestion in enumerate(suggestions, 1):
             result += f"{i}. {suggestion}\n\n"
 
-        print(f"Final result: {result}")
+        print(f"Final result: {result[:100]}...")
 
     # If suggestion type is "common_phrases", use the person's common phrases
     elif suggestion_type == "common_phrases":
@@ -194,12 +248,16 @@ def generate_suggestions(person_id, user_input, suggestion_type, selected_topic=
             print("No category inferred, falling back to model")
             # Fall back to model if we couldn't infer a category
             try:
-                suggestion = suggestion_generator.generate_suggestion(
-                    person_context, user_input
-                )
-                print(f"Generated fallback suggestion: {suggestion}")
-                result = "### AI-Generated Response (no category detected):\n\n"
-                result += f"1. {suggestion}\n\n"
+                suggestions = []
+                for i in range(3):
+                    suggestion = suggestion_generator.generate_suggestion(
+                        person_context, user_input, temperature=temperature
+                    )
+                    suggestions.append(suggestion)
+
+                result = f"### AI-Generated Responses (no category detected, using {suggestion_generator.model_name}):\n\n"
+                for i, suggestion in enumerate(suggestions, 1):
+                    result += f"{i}. {suggestion}\n\n"
             except Exception as e:
                 print(f"Error generating fallback suggestion: {e}")
                 result = "### Could not generate a response:\n\n"
@@ -319,9 +377,33 @@ with gr.Blocks(title="Will's AAC Communication Aid") as demo:
                 info="Choose what kind of responses you want (model = AI-generated)",
             )
 
+            # Model selection
+            with gr.Row():
+                model_dropdown = gr.Dropdown(
+                    choices=list(AVAILABLE_MODELS.keys()),
+                    value="distilgpt2",
+                    label="Language Model",
+                    info="Select which AI model to use for generating responses",
+                )
+
+                temperature_slider = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.5,
+                    value=0.7,
+                    step=0.1,
+                    label="Temperature",
+                    info="Controls randomness (higher = more creative, lower = more focused)",
+                )
+
             # Generate button
             generate_btn = gr.Button("Generate My Responses", variant="primary")
 
+            # Model status
+            model_status = gr.Markdown(
+                value=f"Current model: {suggestion_generator.model_name}",
+                label="Model Status",
+            )
+
         with gr.Column(scale=1):
             # Common phrases
             common_phrases = gr.Textbox(
@@ -347,6 +429,11 @@ with gr.Blocks(title="Will's AAC Communication Aid") as demo:
         # Update the context, phrases, and topic dropdown
         return context_info, phrases_text, gr.update(choices=topics)
 
+    def handle_model_change(model_name):
+        """Handle model selection change."""
+        status = change_model(model_name)
+        return status
+
     # Set up the person change event
     person_dropdown.change(
         handle_person_change,
@@ -354,10 +441,24 @@ with gr.Blocks(title="Will's AAC Communication Aid") as demo:
         outputs=[context_display, common_phrases, topic_dropdown],
     )
 
+    # Set up the model change event
+    model_dropdown.change(
+        handle_model_change,
+        inputs=[model_dropdown],
+        outputs=[model_status],
+    )
+
     # Set up the generate button click event
     generate_btn.click(
         generate_suggestions,
-        inputs=[person_dropdown, user_input, suggestion_type, topic_dropdown],
+        inputs=[
+            person_dropdown,
+            user_input,
+            suggestion_type,
+            topic_dropdown,
+            model_dropdown,
+            temperature_slider,
+        ],
         outputs=[suggestions_output],
     )
 

demo.py

@@ -398,7 +398,10 @@ class LLMToolInterface(LLMInterface):
                         print("llm install llm-mlx")
                     elif "ollama" in self.model_name.lower():
                         print("llm install llm-ollama")
-                        print("ollama pull " + self.model_name.replace("ollama/", ""))
+                        model_name = self.model_name
+                        if "/" in model_name:
+                            model_name = model_name.split("/")[1]
+                        print("ollama pull " + model_name)
             else:
                 print("Warning: LLM tool may be installed but returned an error.")
         except Exception as e:
@@ -610,7 +613,7 @@ def main():
         "- hf: 'distilgpt2', 'gpt2-medium', 'google/gemma-2b-it'\n"
         "- llm: 'gemini-1.5-pro-latest', 'gemma-3-27b-it' (requires llm-gemini plugin)\n"
         "       'mlx-community/gemma-7b-it' (requires llm-mlx plugin)\n"
-        "       'ollama/gemma3:4b-it-qat', 'ollama/llama3:8b' (requires llm-ollama plugin)",
+        "       'Ollama: gemma3:4b-it-qat', 'Ollama: llama3:8b' (requires llm-ollama plugin)",
     )
     parser.add_argument(
         "--num_responses", type=int, default=3, help="Number of responses to generate"
@@ -705,9 +708,12 @@ def main():
                 print("1. Install from https://ollama.ai/")
                 print("2. Start Ollama with: ollama serve")
                 print("3. Install the llm-ollama plugin: llm install llm-ollama")
-                print(
-                    f"4. Pull the model: ollama pull {args.model.replace('ollama/', '')}"
-                )
+                model_name = args.model
+                if "ollama:" in model_name.lower():
+                    model_name = model_name.replace("Ollama: ", "")
+                elif "/" in model_name:
+                    model_name = model_name.split("/")[1]
+                print(f"4. Pull the model: ollama pull {model_name}")
             else:
                 print("\nMake sure Simon Willison's LLM tool is installed:")
                 print("pip install llm")

utils.py

@@ -159,16 +159,20 @@ class SuggestionGenerator:
         """
         self.model_name = model_name
         self.model_loaded = False
+        self.generator = None
+        self.aac_user_info = None
 
+        # Load AAC user information from social graph
         try:
-            print(f"Loading model: {model_name}")
-            # Use a simpler approach with a pre-built pipeline
-            self.generator = pipeline("text-generation", model=model_name)
-            self.model_loaded = True
-            print(f"Model loaded successfully: {model_name}")
+            with open("social_graph.json", "r") as f:
+                social_graph = json.load(f)
+                self.aac_user_info = social_graph.get("aac_user", {})
         except Exception as e:
-            print(f"Error loading model: {e}")
-            self.model_loaded = False
+            print(f"Error loading AAC user info from social graph: {e}")
+            self.aac_user_info = {}
+
+        # Try to load the model
+        self.load_model(model_name)
 
         # Fallback responses if model fails to load or generate
         self.fallback_responses = [
@@ -176,8 +180,92 @@ class SuggestionGenerator:
             "That's interesting. Tell me more.",
             "I'd like to talk about that further.",
             "I appreciate you sharing that with me.",
+            "Could we talk about something else?",
+            "I need some time to think about that.",
         ]
 
+    def load_model(self, model_name: str) -> bool:
+        """Load a Hugging Face model.
+
+        Args:
+            model_name: Name of the HuggingFace model to use
+
+        Returns:
+            bool: True if model loaded successfully, False otherwise
+        """
+        self.model_name = model_name
+        self.model_loaded = False
+
+        try:
+            print(f"Loading model: {model_name}")
+
+            # Check if this is a gated model that requires authentication
+            is_gated_model = any(
+                name in model_name.lower()
+                for name in ["gemma", "llama", "mistral", "qwen", "phi"]
+            )
+
+            if is_gated_model:
+                # Try to get token from environment
+                import os
+
+                token = os.environ.get("HUGGING_FACE_HUB_TOKEN") or os.environ.get(
+                    "HF_TOKEN"
+                )
+
+                if token:
+                    print(f"Using token for gated model: {model_name}")
+                    from huggingface_hub import login
+
+                    login(token=token, add_to_git_credential=False)
+
+                    # Explicitly pass token to pipeline
+                    from transformers import AutoTokenizer, AutoModelForCausalLM
+
+                    try:
+                        tokenizer = AutoTokenizer.from_pretrained(
+                            model_name, token=token
+                        )
+                        model = AutoModelForCausalLM.from_pretrained(
+                            model_name, token=token
+                        )
+                        self.generator = pipeline(
+                            "text-generation", model=model, tokenizer=tokenizer
+                        )
+                    except Exception as e:
+                        print(f"Error loading gated model with token: {e}")
+                        print(
+                            "This may be due to not having accepted the model license or insufficient permissions."
+                        )
+                        print(
+                            "Please visit the model page on Hugging Face Hub and accept the license."
+                        )
+                        raise
+                else:
+                    print("No Hugging Face token found in environment variables.")
+                    print(
+                        "To use gated models like Gemma, you need to set up a token with the right permissions."
+                    )
+                    print("1. Create a token at https://huggingface.co/settings/tokens")
+                    print(
+                        "2. Make sure to enable 'Access to public gated repositories'"
+                    )
+                    print(
+                        "3. Set it as an environment variable: export HUGGING_FACE_HUB_TOKEN=your_token_here"
+                    )
+                    raise ValueError("Authentication token required for gated model")
+            else:
+                # For non-gated models, use the standard pipeline
+                self.generator = pipeline("text-generation", model=model_name)
+
+            self.model_loaded = True
+            print(f"Model loaded successfully: {model_name}")
+            return True
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            self.model_loaded = False
+            return False
+
     def test_model(self) -> str:
         """Test if the model is working correctly."""
         if not self.model_loaded:
@@ -186,7 +274,9 @@ class SuggestionGenerator:
         try:
             test_prompt = "I am Will. My son Billy asked about football. I respond:"
             print(f"Testing model with prompt: {test_prompt}")
-            response = self.generator(test_prompt, max_length=30, do_sample=True)
+            response = self.generator(
+                test_prompt, max_new_tokens=30, do_sample=True, truncation=True
+            )
             result = response[0]["generated_text"][len(test_prompt) :]
             print(f"Test response: {result}")
             return f"Model test successful: {result}"
@@ -222,39 +312,100 @@ class SuggestionGenerator:
         # Extract context information
         name = person_context.get("name", "")
         role = person_context.get("role", "")
-        topics = ", ".join(person_context.get("topics", []))
+        topics = person_context.get("topics", [])
         context = person_context.get("context", "")
         selected_topic = person_context.get("selected_topic", "")
+        common_phrases = person_context.get("common_phrases", [])
+        frequency = person_context.get("frequency", "")
 
-        # Build prompt
-        prompt = f"""I am Will, a person with MND (Motor Neuron Disease).
-I'm talking to {name}, who is my {role}.
-"""
+        # Get AAC user information
+        aac_user = self.aac_user_info
 
-        if context:
-            prompt += f"Context: {context}\n"
+        # Build enhanced prompt
+        prompt = f"""I am {aac_user.get('name', 'Will')}, a {aac_user.get('age', 38)}-year-old with MND (Motor Neuron Disease) from {aac_user.get('location', 'Manchester')}.
+{aac_user.get('background', '')}
 
-        if topics:
-            prompt += f"Topics of interest: {topics}\n"
+My communication needs: {aac_user.get('communication_needs', '')}
 
-        if selected_topic:
-            prompt += f"We're currently talking about: {selected_topic}\n"
+I am talking to {name}, who is my {role}.
+About {name}: {context}
+We typically talk about: {', '.join(topics)}
+We communicate {frequency}.
+"""
+
+        # Add communication style based on relationship
+        if role in ["wife", "son", "daughter", "mother", "father"]:
+            prompt += "I communicate with my family in a warm, loving way, sometimes using inside jokes.\n"
+        elif role in ["doctor", "therapist", "nurse"]:
+            prompt += "I communicate with healthcare providers in a direct, informative way.\n"
+        elif role in ["best mate", "friend"]:
+            prompt += "I communicate with friends casually, often with humor and sometimes swearing.\n"
+        elif role in ["work colleague", "boss"]:
+            prompt += (
+                "I communicate with colleagues professionally but still friendly.\n"
+            )
 
+        # Add topic information if provided
+        if selected_topic:
+            prompt += f"\nWe are currently discussing {selected_topic}.\n"
+
+            # Add specific context about this topic with this person
+            if selected_topic == "football" and "Manchester United" in context:
+                prompt += "We both support Manchester United and often discuss recent matches.\n"
+            elif selected_topic == "programming" and "software developer" in context:
+                prompt += "We both work in software development and share technical interests.\n"
+            elif selected_topic == "family plans" and role in ["wife", "husband"]:
+                prompt += (
+                    "We make family decisions together, considering my condition.\n"
+                )
+            elif selected_topic == "old scout adventures" and role == "best mate":
+                prompt += "We often reminisce about our Scout camping trips in South East London.\n"
+            elif selected_topic == "cycling" and "cycling" in context:
+                prompt += "I miss being able to cycle but enjoy talking about past cycling adventures.\n"
+
+        # Add the user's message if provided
         if user_input:
             prompt += f'\n{name} just said to me: "{user_input}"\n'
-
-        prompt += "\nMy response:"
+        elif common_phrases:
+            # Use a common phrase from the person if no message is provided
+            default_message = common_phrases[0]
+            prompt += f'\n{name} typically says things like: "{default_message}"\n'
+
+        # Add the response prompt with specific guidance
+        # Check if this is an instruction-tuned model
+        is_instruction_model = any(
+            marker in self.model_name.lower()
+            for marker in ["-it", "instruct", "chat", "phi-3", "phi-2"]
+        )
+
+        if is_instruction_model:
+            # Use instruction format for instruction-tuned models
+            prompt += f"""
+<instruction>
+Respond to {name} in a way that is natural, brief (1-2 sentences), and directly relevant to what they just said.
+Use language appropriate for our relationship.
+</instruction>
+
+My response to {name}:"""
+        else:
+            # Use standard format for non-instruction models
+            prompt += f"""
+I want to respond to {name} in a way that is natural, brief (1-2 sentences), and directly relevant to what they just said. I'll use language appropriate for our relationship.
+
+My response to {name}:"""
 
         # Generate suggestion
         try:
             print(f"Generating suggestion with prompt: {prompt}")
+            # Use max_new_tokens instead of max_length to avoid the error
             response = self.generator(
                 prompt,
-                max_length=len(prompt.split()) + max_length,
+                max_new_tokens=max_length,  # Generate new tokens, not including prompt
                 temperature=temperature,
                 do_sample=True,
                 top_p=0.92,
                 top_k=50,
+                truncation=True,
             )
             # Extract only the generated part, not the prompt
             result = response[0]["generated_text"][len(prompt) :]