Update app.py

app.py

@@ -1,156 +1,152 @@
 import gradio as gr
-import random
+from typing import List, Tuple, Optional, Union
 import os
-
-# Attempt to import required libraries
-try:
-    from diffusers import StableDiffusionPipeline
-    import torch
-    STABLE_DIFFUSION_AVAILABLE = True
-except ImportError as e:
-    print(f"Error importing Stable Diffusion dependencies: {e}")
-    STABLE_DIFFUSION_AVAILABLE = False
-
-try:
-    import openai
-    openai.api_key = os.environ.get("OPENAI_API_KEY")
-    USE_GPT = True
-except ImportError:
-    print("OpenAI library not found. Falling back to basic responses.")
-    USE_GPT = False
-
-# Quiz questions and answers
-christmas_quiz = [
-    {
-        "question": "What is the traditional Christmas flower?",
-        "options": ["Rose", "Poinsettia", "Tulip", "Daisy"],
-        "answer": "Poinsettia"
-    },
-    {
-        "question": "In which country did the tradition of putting up a Christmas tree originate?",
-        "options": ["USA", "England", "Germany", "France"],
-        "answer": "Germany"
-    },
-    {
-        "question": "What is the name of the ballet often performed at Christmas?",
-        "options": ["Swan Lake", "The Nutcracker", "Sleeping Beauty", "Giselle"],
-        "answer": "The Nutcracker"
-    },
-    {
-        "question": "Which company was the first to use Santa Claus in advertising?",
-        "options": ["Pepsi", "Coca-Cola", "McDonald's", "Walmart"],
-        "answer": "Coca-Cola"
-    },
-    {
-        "question": "What is the most popular Christmas dinner in Japan?",
-        "options": ["Turkey", "Ham", "KFC Chicken", "Roast Beef"],
-        "answer": "KFC Chicken"
-    }
-]
-
-# Initialize the Stable Diffusion pipeline if available
-if STABLE_DIFFUSION_AVAILABLE:
-    model_id = "runwayml/stable-diffusion-v1-5"
-    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
-
-def generate_image(prompt):
-    if not STABLE_DIFFUSION_AVAILABLE:
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+class ChristmasBot:
+    def __init__(self, model_name: str = "TheBloke/Mistral-7B-v0.1-GGUF"):
+        """
+        Initialize the Christmas chatbot with a Hugging Face model.
+        Default model is Mistral-7B, but you can change it to any other model.
+        """
+        self.stable_diffusion_available = False
+        
+        # Initialize the model and tokenizer
+        print("Loading model and tokenizer...")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            load_in_8bit=True  # Enable 8-bit quantization for memory efficiency
+        )
+        print("Model loaded successfully!")
+        
+        # System prompt to give the model Christmas context
+        self.system_prompt = """You are Holly, a cheerful Christmas helper chatbot. 
+        You love everything about Christmas and respond in a warm, festive manner. 
+        Keep your responses concise but friendly. 
+        If users ask about sensitive topics, guide the conversation back to Christmas-related subjects."""
+        
+    def _generate_image(self, prompt: str) -> Optional[str]:
+        """Placeholder for image generation functionality."""
+        if not self.stable_diffusion_available:
+            return None
         return None
-    with torch.no_grad():
-        image = pipe(prompt, num_inference_steps=50).images[0]
-    return image
-
-def get_gpt_response(prompt, history):
-    if not USE_GPT:
-        return "I'm sorry, but I'm currently operating with limited capabilities. I can still help with the Christmas quiz and card generation!"
     
-    messages = [
-        {"role": "system", "content": "You are a helpful Christmas-themed chatbot named Holly. You can answer questions about Christmas, offer holiday tips, and engage in festive conversation. You also know about the Christmas quiz and card generation features of this application."},
-    ]
-    for h in history:
-        messages.append({"role": "user", "content": h[0]})
-        messages.append({"role": "assistant", "content": h[1]})
-    messages.append({"role": "user", "content": prompt})
+    def _format_chat_history(self, history: List[List[str]]) -> str:
+        """Format the chat history into a single string for the model."""
+        formatted_history = self.system_prompt + "\n\n"
+        for user_msg, bot_msg in history:
+            if user_msg:
+                formatted_history += f"User: {user_msg}\n"
+            if bot_msg:
+                formatted_history += f"Assistant: {bot_msg}\n"
+        return formatted_history
     
-    try:
-        response = openai.ChatCompletion.create(
-            model="gpt-3.5-turbo",
-            messages=messages,
-            max_tokens=150,
-            n=1,
-            stop=None,
-            temperature=0.7,
-        )
-        return response.choices[0].message['content'].strip()
-    except Exception as e:
-        print(f"Error in GPT response: {e}")
-        return "I'm having trouble connecting to my knowledge base right now. Can I help you with the Christmas quiz or card generation instead?"
-
-def chatbot(message, history):
-    if not history:
-        return "Ho ho ho! Merry Christmas! I'm Holly, your Christmas helper. Would you like to take a Christmas quiz, create a Christmas card, or chat about the holidays?"
-
-    last_response = history[-1][1].lower()
-
-    if "quiz" in message.lower():
-        question = random.choice(christmas_quiz)
-        options_text = "\n".join([f"{i+1}. {opt}" for i, opt in enumerate(question['options'])])
-        return f"Great! Here's your Christmas quiz question:\n\n{question['question']}\n\n{options_text}\n\nPlease enter the number of your answer."
-
-    elif "card" in message.lower():
-        if STABLE_DIFFUSION_AVAILABLE:
-            return "Wonderful! Let's create a Christmas card. Please describe the scene you'd like on your card, and I'll generate it for you using AI."
-        else:
-            return "I'm sorry, but the card generation feature is currently unavailable. Would you like to take a Christmas quiz instead?"
-
-    elif any(str(i) in message for i in range(1, 5)):  # Check if the message is a quiz answer
-        for q in christmas_quiz:
-            if q['question'] in history[-2][1]:  # Find the question in the history
-                user_answer = q['options'][int(message) - 1]
-                if user_answer == q['answer']:
-                    return f"Correct! {q['answer']} is the right answer. Would you like another question, to create a Christmas card, or to chat about something else?"
-                else:
-                    return f"Sorry, that's not correct. The right answer is {q['answer']}. Would you like another question, to create a Christmas card, or to chat about something else?"
-
-    elif "card" in last_response and STABLE_DIFFUSION_AVAILABLE:
-        image = generate_image(f"Christmas card scene: {message}")
-        if image:
-            return (f"I've created a Christmas card based on your description: '{message}'. You can see it in the image box below. "
-                    f"Would you like to create another card, take a quiz, or chat about something else?", image)
-        else:
-            return "I'm sorry, I couldn't generate the image. Would you like to try again, take a quiz, or chat about something else?"
-
-    else:
-        # Use GPT-3.5 for general conversation
-        return get_gpt_response(message, history)
-
-# Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# Christmas Quiz and Card Generator Chatbot")
-    gr.Markdown("""
-    Welcome to the Christmas Quiz and Card Generator Chatbot!
-    - Type 'quiz' to start a Christmas quiz.
-    - Type 'card' to create a custom Christmas card.
-    - Or just chat about anything Christmas-related!
-    """)
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Type your message here")
-    clear = gr.Button("Clear")
-
-    def user(user_message, history):
+    def _get_llm_response(self, message: str, history: List[List[str]]) -> str:
+        """Generate response using the Hugging Face model."""
+        try:
+            # Format the conversation history with the new message
+            chat_history = self._format_chat_history(history)
+            prompt = f"{chat_history}User: {message}\nAssistant:"
+            
+            # Tokenize input
+            inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
+            inputs = inputs.to(self.model.device)
+            
+            # Generate response
+            outputs = self.model.generate(
+                inputs["input_ids"],
+                max_length=2048,
+                temperature=0.7,
+                top_p=0.9,
+                repetition_penalty=1.2,
+                do_sample=True,
+                num_return_sequences=1,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
+            
+            # Decode and clean up the response
+            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            response = response.split("Assistant:")[-1].strip()
+            
+            return response
+            
+        except Exception as e:
+            print(f"Error generating response: {e}")
+            return "Ho ho ho! I seem to be having trouble with my Christmas magic. Could you try asking me something else?"
+    
+    def process_message(self, message: str, history: List[List[str]]) -> Union[str, Tuple[str, str]]:
+        """Process user message and return appropriate response."""
+        # Initial greeting
+        if not history:
+            return "Ho ho ho! Merry Christmas! I'm Holly, your Christmas helper. Would you like to create a Christmas card or chat about the holidays?"
+        
+        message = message.lower()
+        last_response = history[-1][1].lower() if history else ""
+        
+        # Handle card creation request
+        if "card" in message:
+            if self.stable_diffusion_available:
+                return "Wonderful! Let's create a Christmas card. Please describe the scene you'd like on your card, and I'll generate it for you using AI."
+            return "I'm sorry, but the card generation feature is currently unavailable. Let's chat about Christmas instead!"
+        
+        # Handle card generation
+        if "card" in last_response and self.stable_diffusion_available:
+            image = self._generate_image(f"Christmas card scene: {message}")
+            if image:
+                return (
+                    f"I've created a Christmas card based on your description: '{message}'. "
+                    f"Would you like to create another card or chat about something else?",
+                    image
+                )
+            return "I'm sorry, I couldn't generate the image. Would you like to try again or chat about something else?"
+        
+        # Default to LLM response for all other messages
+        return self._get_llm_response(message, history)
+
+def create_gradio_interface() -> gr.Blocks:
+    """Create and configure the Gradio interface."""
+    # You can change the model here
+    bot = ChristmasBot(model_name="TheBloke/Mistral-7B-v0.1-GGUF")
+    
+    def user(user_message: str, history: List[List[str]]) -> Tuple[str, List[List[str]]]:
         return "", history + [[user_message, None]]
-
-    def bot(history):
-        bot_message = chatbot(history[-1][0], history[:-1])
-        history[-1][1] = bot_message
+    
+    def bot_response(history: List[List[str]]) -> Tuple[List[List[str]], Optional[str]]:
+        bot_message = bot.process_message(history[-1][0], history[:-1])
         if isinstance(bot_message, tuple):
+            history[-1][1] = bot_message[0]
             return history, bot_message[1]
+        history[-1][1] = bot_message
         return history, None
+    
+    with gr.Blocks() as demo:
+        gr.Markdown("# 🎄 Christmas Chatbot & Card Generator 🎅")
+        gr.Markdown("""
+        Welcome to the Christmas Chatbot!
+        - Chat about anything Christmas-related
+        - Type 'card' to create a custom Christmas card
+        """)
+        
+        chatbot = gr.Chatbot()
+        msg = gr.Textbox(
+            label="Type your message here",
+            placeholder="Ask me anything about Christmas or request a card!",
+            show_label=True
+        )
+        clear = gr.Button("Clear Chat")
+        image_output = gr.Image(label="Generated Card")
+        
+        msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+            bot_response, chatbot, [chatbot, image_output]
+        )
+        clear.click(lambda: None, None, chatbot, queue=False)
+    
+    return demo
 
-    image_output = gr.Image()
-    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, chatbot, [chatbot, image_output]
-    )
-    clear.click(lambda: None, None, chatbot, queue=False)
-
-demo.launch()
+if __name__ == "__main__":
+    demo = create_gradio_interface()
+    demo.launch()