Update app.py

app.py

@@ -1,101 +1,77 @@
 import gradio as gr
 from typing import List, Tuple, Optional, Union
 import os
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+from openai import OpenAI
+import json
 
 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.
-        """
+    def __init__(self):
+        """Initialize the Christmas chatbot with OpenAI."""
         self.stable_diffusion_available = False
+        # For Hugging Face Spaces, get the API key from secrets
+        self.client = OpenAI(api_key=os.environ.get('OPENAI_API_KEY'))
         
-        # 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:
+        """Generate an image using DALL-E."""
+        try:
+            response = self.client.images.generate(
+                model="dall-e-3",
+                prompt=f"Christmas themed illustration: {prompt}, festive, cheerful, holiday spirit",
+                size="1024x1024",
+                quality="standard",
+                n=1,
+            )
+            return response.data[0].url
+        except Exception as e:
+            print(f"Image generation error: {e}")
             return None
-        return None
-    
-    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
     
     def _get_llm_response(self, message: str, history: List[List[str]]) -> str:
-        """Generate response using the Hugging Face model."""
+        """Get response from OpenAI."""
         try:
-            # Format the conversation history with the new message
-            chat_history = self._format_chat_history(history)
-            prompt = f"{chat_history}User: {message}\nAssistant:"
+            messages = [{"role": "system", "content": self.system_prompt}]
             
-            # Tokenize input
-            inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
-            inputs = inputs.to(self.model.device)
+            for h in history:
+                if h[0]:
+                    messages.append({"role": "user", "content": h[0]})
+                if h[1]:
+                    messages.append({"role": "assistant", "content": h[1]})
             
-            # Generate response
-            outputs = self.model.generate(
-                inputs["input_ids"],
-                max_length=2048,
+            messages.append({"role": "user", "content": message})
+            
+            response = self.client.chat.completions.create(
+                model="gpt-3.5-turbo",
+                messages=messages,
                 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
+                max_tokens=150
             )
             
-            # Decode and clean up the response
-            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            response = response.split("Assistant:")[-1].strip()
-            
-            return response
+            return response.choices[0].message.content
             
         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?"
+            print(f"OpenAI API error: {e}")
+            return "Ho ho ho! I seem to be having trouble with my Christmas magic. Could you try again?"
     
     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 "Wonderful! Let's create a Christmas card. Please describe the scene you'd like on your card, and I'll generate it using DALL-E."
             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}")
+            image = self._generate_image(message)
             if image:
                 return (
                     f"I've created a Christmas card based on your description: '{message}'. "
@@ -104,13 +80,61 @@ class ChristmasBot:
                 )
             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")
+with gr.Blocks(css="""
+    :root {
+        --holly-green: #146B3A;
+        --christmas-red: #EA4630;
+        --snow-white: #F8F9FA;
+        --gold: #FFC107;
+    }
+    .message.user {
+        background-color: var(--holly-green) !important;
+        color: white !important;
+        border-radius: 15px 15px 5px 15px !important;
+    }
+    .message.bot {
+        background-color: var(--christmas-red) !important;
+        color: white !important;
+        border-radius: 15px 15px 15px 5px !important;
+    }
+    textarea {
+        border: 2px solid var(--holly-green) !important;
+        border-radius: 8px !important;
+    }
+    button {
+        background-color: var(--holly-green) !important;
+        color: white !important;
+        border: none !important;
+        border-radius: 8px !important;
+        transition: background-color 0.2s !important;
+    }
+    button:hover {
+        background-color: var(--christmas-red) !important;
+    }
+""") 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 with DALL-E
+    """)
+    
+    bot = ChristmasBot()
+    
+    chatbot = gr.Chatbot(
+        bubble_full_width=False,
+        avatar_images=("👤", "🎅"),
+        height=400
+    )
+    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", visible=False)
     
     def user(user_message: str, history: List[List[str]]) -> Tuple[str, List[List[str]]]:
         return "", history + [[user_message, None]]
@@ -123,30 +147,10 @@ def create_gradio_interface() -> gr.Blocks:
         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
+    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot_response, chatbot, [chatbot, image_output]
+    )
+    clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":
-    demo = create_gradio_interface()
     demo.launch()