Update app.py

app.py

@@ -1,110 +1,80 @@
 import gradio as gr
-from huggingface_hub import InferenceClient, InferenceTimeoutError
-import logging
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
+# Load your fine-tuned model and tokenizer
+model = AutoModelForCausalLM.from_pretrained(
+    "hackergeek/gemma-finetuned",
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
+tokenizer = AutoTokenizer.from_pretrained("hackergeek/gemma-finetuned")
+tokenizer.pad_token = tokenizer.eos_token
 
-# Initialize client with timeout
-client = InferenceClient("hackergeek/gemma-finetuned", timeout=30)
-
-def respond(
-    message: str,
-    history: list[tuple[str, str]],
-    system_message: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
-):
-    """
-    Handles chat responses with error handling and validation
-    """
-    try:
-        # Validate system message
-        if not system_message.strip():
-            system_message = "You are a helpful AI assistant."
-
-        # Build message history
-        messages = [{"role": "system", "content": system_message}]
-        
-        for user_msg, assistant_msg in history:
-            if user_msg:
-                messages.append({"role": "user", "content": user_msg})
-            if assistant_msg:
-                messages.append({"role": "assistant", "content": assistant_msg})
-        
-        messages.append({"role": "user", "content": message})
-
-        response = ""
-        
-        # Stream the response
-        for chunk in client.chat_completion(
-            messages,
-            max_tokens=max_tokens,
-            stream=True,
-            temperature=temperature,
-            top_p=top_p,
-        ):
-            if chunk.choices and chunk.choices[0].delta.content:
-                token = chunk.choices[0].delta.content
-                response += token
-                yield response
+def format_prompt(message, history):
+    """Format the prompt with conversation history"""
+    system_prompt = "You are a knowledgeable space expert assistant. Answer questions about astronomy, space exploration, and related topics in a clear and engaging manner."
+    prompt = f"<system>{system_prompt}</system>\n"
+    
+    for user_msg, bot_msg in history:
+        prompt += f"<user>{user_msg}</user>\n<assistant>{bot_msg}</assistant>\n"
+    
+    prompt += f"<user>{message}</user>\n<assistant>"
+    return prompt
 
-    except InferenceTimeoutError:
-        logging.error("API request timed out")
-        yield "Error: Request timed out. Please try again."
-    except Exception as e:
-        logging.error(f"API error: {str(e)}")
-        yield f"Error: {str(e)}. Please check your input and try again."
+def respond(message, history):
+    # Format the prompt with conversation history
+    full_prompt = format_prompt(message, history)
+    
+    # Tokenize input
+    inputs = tokenizer(full_prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
+    
+    # Generate response
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=1024,
+        temperature=0.7,
+        top_p=0.9,
+        repetition_penalty=1.1,
+        do_sample=True
+    )
+    
+    # Decode and extract only the new response
+    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+    
+    return response
 
-# Custom CSS for better appearance
-custom_css = """
-#chatbot { min-height: 400px; }
-.dark #chatbot { font-size: 14px !important; }
-footer { visibility: hidden; }
+# Custom CSS for space theme
+space_css = """
+.gradio-container {
+    background: linear-gradient(45deg, #000000, #1a1a2e);
+    color: white;
+}
+.chatbot {
+    background-color: rgba(0, 0, 0, 0.7) !important;
+    border: 1px solid #4a4a4a !important;
+}
 """
 
-# Configure interface
-with gr.Blocks(css=custom_css, title="Gemma Chatbot") as demo:
-    gr.Markdown("# 🚀 Gemma Fine-Tuned Chatbot")
-    gr.Markdown("Chat with the fine-tuned Gemma AI assistant!")
+# Create the interface
+with gr.Blocks(css=space_css, theme=gr.themes.Default(primary_hue="blue", secondary_hue="purple")) as demo:
+    gr.Markdown("# 🚀 Space Explorer Chatbot 🌌")
+    gr.Markdown("Ask me anything about space! Planets, stars, galaxies, or space exploration!")
     
-    with gr.Accordion("⚙️ Advanced Settings", open=False):
-        system_input = gr.Textbox(
-            value="You are a helpful AI assistant.",
-            label="System Role",
-            info="Initial instructions for the AI"
-        )
-        max_tokens = gr.Slider(
-            minimum=32, maximum=2048, value=512,
-            step=32, label="Max Response Length"
-        )
-        temperature = gr.Slider(
-            minimum=0.1, maximum=2.0, value=0.7,
-            step=0.1, label="Creativity (Temperature)"
-        )
-        top_p = gr.Slider(
-            minimum=0.1, maximum=1.0, value=0.95,
-            step=0.05, label="Focus (Top-p)"
-        )
-    
-    chat_interface = gr.ChatInterface(
+    chatbot = gr.ChatInterface(
         respond,
-        additional_inputs=[
-            system_input,
-            max_tokens,
-            temperature,
-            top_p
-        ],
         examples=[
-            ["Explain quantum computing in simple terms"],
-            ["What's the weather like in Paris?"],
-            ["Write a poem about artificial intelligence"]
+            "Explain black holes in simple terms",
+            "What's the latest news about Mars exploration?",
+            "How do stars form?",
+            "Tell me about the James Webb Space Telescope"
         ],
         retry_btn=None,
         undo_btn=None,
-        clear_btn="✨ New Chat",
+        clear_btn="Clear History",
     )
+    
+    chatbot.chatbot.height = 600
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(share=True)