Update app.py

app.py

@@ -6,22 +6,43 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 # Using an open-access model instead of gated Mistral
 MODEL_NAME = "tiiuae/falcon-7b-instruct"
 
-# Preload model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto"
-)
-
-generator = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=512,
-    temperature=0.5,
-    do_sample=True
-)
+try:
+    # Preload model and tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    
+    # Load model with appropriate settings
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto" if torch.cuda.is_available() else None,
+        low_cpu_mem_usage=True
+    )
+    
+    generator = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=512,
+        temperature=0.5,
+        do_sample=True
+    )
+    
+except Exception as e:
+    print(f"Error loading model: {str(e)}")
+    # Fallback to CPU if GPU fails
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float32,
+        device_map=None
+    )
+    generator = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=512,
+        temperature=0.5,
+        do_sample=True
+    )
 
 # ---------- TECH FILTER ----------
 def is_tech_query(message: str) -> bool:
@@ -29,59 +50,93 @@ def is_tech_query(message: str) -> bool:
         "python", "java", "javascript", "html", "css", "react", "angular",
         "node", "machine learning", "deep learning", "ai", "api", "code",
         "debug", "error", "technology", "computer", "programming", "software",
-        "hardware", "cybersecurity", "database", "sql", "devops", "cloud"
+        "hardware", "cybersecurity", "database", "sql", "devops", "cloud",
+        "algorithm", "backend", "frontend", "server", "linux", "windows",
+        "docker", "kubernetes", "git", "github", "vscode", "pycharm",
+        "tensorflow", "pytorch", "neural network", "blockchain", "web3",
+        "smart contract", "ethereum", "bitcoin", "cryptography", "encryption"
     ]
     return any(k in message.lower() for k in tech_keywords)
 
 # ---------- CHAT FUNCTION ----------
 def chat_with_model(message, history):
     if not is_tech_query(message):
-        return history + [[message, "⚠️ I can only answer technology-related queries."]]
+        return history + [[message, "⚠️ I can only answer technology-related queries. Please ask about programming, AI, cybersecurity, or other tech topics."]]
 
     conversation = ""
     for user_msg, bot_msg in history:
         conversation += f"User: {user_msg}\nAssistant: {bot_msg}\n"
     conversation += f"User: {message}\nAssistant:"
 
-    output = generator(conversation)[0]["generated_text"]
-    if "Assistant:" in output:
-        answer = output.split("Assistant:")[-1].strip()
-    else:
-        answer = output.strip()
-
-    return history + [[message, answer]]
+    try:
+        output = generator(conversation, pad_token_id=tokenizer.eos_token_id)[0]["generated_text"]
+        if "Assistant:" in output:
+            answer = output.split("Assistant:")[-1].strip()
+        else:
+            answer = output.strip()
+        
+        # Clean up response
+        answer = answer.split("User:")[0].strip()  # Remove any following user prompts
+        return history + [[message, answer]]
+    
+    except Exception as e:
+        error_msg = f"❌ Error generating response: {str(e)}"
+        return history + [[message, error_msg]]
 
 # ---------- LOGIN + UI ----------
 session_state = {"authenticated": False}
 
 def login(username, password):
-    if (username == "admin" and password == "admin123") or (username == "techuser" and password == "techpass"):
+    valid_credentials = {
+        "admin": "admin123",
+        "techuser": "techpass",
+        "guest": "guest123"
+    }
+    
+    if username in valid_credentials and password == valid_credentials[username]:
         session_state["authenticated"] = True
         return gr.update(visible=False), gr.update(visible=True), ""
     else:
-        return gr.update(), gr.update(visible=False), "❌ Invalid credentials."
+        return gr.update(), gr.update(visible=False), "❌ Invalid credentials. Try admin/admin123 or techuser/techpass"
+
+def logout():
+    session_state["authenticated"] = False
+    return gr.update(visible=True), gr.update(visible=False), "Logged out successfully"
 
-with gr.Blocks(css=".gradio-container {max-width: 750px; margin: auto;}") as demo:
+with gr.Blocks(css="""
+    .gradio-container {max-width: 750px; margin: auto;}
+    .chatbot {min-height: 500px;}
+    """) as demo:
+    
     # Login Page
-    with gr.Group(visible=not session_state["authenticated"]) as login_group:
-        gr.Markdown("# 🔐 Login to Tech Chatbot")
-        username = gr.Textbox(label="Username")
-        password = gr.Textbox(label="Password", type="password")
-        login_btn = gr.Button("Login")
+    with gr.Group(visible=True) as login_group:
+        gr.Markdown("# 🔐 Tech Chatbot Login")
+        with gr.Row():
+            username = gr.Textbox(label="Username", placeholder="Enter your username")
+            password = gr.Textbox(label="Password", type="password", placeholder="Enter your password")
+        with gr.Row():
+            login_btn = gr.Button("Login", variant="primary")
         login_status = gr.Markdown("")
 
     # Chatbot Page
-    with gr.Group(visible=session_state["authenticated"]) as chat_group:
-        gr.Markdown("# 💻 Tech Helper Chatbot")
+    with gr.Group(visible=False) as chat_group:
+        gr.Markdown("# 💻 Tech Assistant")
         chatbot = gr.Chatbot(height=500)
-        msg = gr.Textbox(placeholder="Type your tech question here...", label="Your Message")
-        clear = gr.Button("Clear Chat")
+        with gr.Row():
+            msg = gr.Textbox(placeholder="Ask about programming, AI, cybersecurity...", 
+                            label="Your Tech Question", scale=4)
+            submit_btn = gr.Button("Send", variant="primary", scale=1)
+        with gr.Row():
+            clear = gr.Button("Clear Chat")
+            logout_btn = gr.Button("Logout")
 
-        msg.submit(chat_with_model, [msg, chatbot], chatbot)
-        clear.click(lambda: None, None, chatbot)
+        msg.submit(chat_with_model, [msg, chatbot], [chatbot])
+        submit_btn.click(chat_with_model, [msg, chatbot], [chatbot])
+        clear.click(lambda: None, None, chatbot, queue=False)
+        logout_btn.click(logout, None, [login_group, chat_group, login_status])
 
     # Button Logic
     login_btn.click(login, [username, password], [login_group, chat_group, login_status])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860)