update space

app.py

@@ -10,6 +10,8 @@ import numpy as np
 from huggingface_hub import login
 from pymongo import MongoClient
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from datetime import datetime
+from werkzeug.security import generate_password_hash, check_password_hash
 
 # Load environment variables from .env file
 dotenv.load_dotenv()
@@ -32,6 +34,8 @@ MONGO_URI = os.getenv("MONGO_URI")
 logger.info("Connecting to MongoDB.")
 client = MongoClient(MONGO_URI)
 db = client.poultry_farm  # Connect to the 'poultry_farm' database
+enquiries_collection = db.enquiries  # Collection to store farmer enquiries
+users_collection = db.users  # Collection to store user credentials
 
 # GPU Setup for TensorFlow
 logger.info("TensorFlow version: %s", tf.__version__)
@@ -74,8 +78,6 @@ def load_model_with_device(model_path, device_name):
 device_name = '/GPU:0' if len(tf.config.list_physical_devices('GPU')) > 0 else '/CPU:0'
 logger.info("Loading disease detection model.")
 my_model = load_model_with_device('models/Final_Chicken_disease_model.h5', device_name)
-logger.info("Loading authentication model.")
-auth_model = load_model_with_device('models/auth_model.h5', device_name)
 
 # Disease names and recommendations
 name_disease = {0: 'Coccidiosis', 1: 'Healthy', 2: 'New Castle Disease', 3: 'Salmonella'}
@@ -172,7 +174,7 @@ class PoultryFarmBot:
             f"Here is some information about {disease_name}: causes, symptoms, and treatment methods "
             "to effectively manage this condition on a poultry farm."
         )
-        response = llama2_response(prompt)
+        response = llama3_response(prompt)
         return response.replace(prompt, "").strip()
 
     def diagnose_disease(self, image):
@@ -191,6 +193,65 @@ class PoultryFarmBot:
         logger.warning("No image provided for diagnosis.")
         return "Please provide an image of poultry fecal matter for disease detection.", None, None, None
 
+    def log_enquiry(self, enquiry_type, content, response, user_id):
+        """
+        Log a farmer's enquiry in the database.
+
+        Args:
+            enquiry_type (str): Type of the enquiry ('image' or 'text').
+            content (str): The content of the enquiry.
+            response (str): The response given by the system.
+            user_id (str): The ID of the user making the enquiry.
+        """
+        enquiry = {
+            "user_id": user_id,
+            "enquiry_type": enquiry_type,
+            "content": content,
+            "response": response,
+            "timestamp": datetime.utcnow()
+        }
+        logger.info(f"Logging enquiry: {enquiry}")
+        enquiries_collection.insert_one(enquiry)
+
+    def authenticate_user(self, username, password):
+        """
+        Authenticate a user with username and password.
+
+        Args:
+            username (str): Username of the user.
+            password (str): Password of the user.
+
+        Returns:
+            dict: User information if authentication is successful, None otherwise.
+        """
+        logger.info(f"Authenticating user: {username}")
+        user = users_collection.find_one({"username": username})
+        if user and check_password_hash(user['password'], password):
+            logger.info("Authentication successful.")
+            return user
+        logger.warning("Authentication failed.")
+        return None
+
+    def register_user(self, username, password):
+        """
+        Register a new user with username and password.
+
+        Args:
+            username (str): Username of the new user.
+            password (str): Password of the new user.
+
+        Returns:
+            bool: True if registration is successful, False otherwise.
+        """
+        logger.info(f"Registering user: {username}")
+        if users_collection.find_one({"username": username}):
+            logger.warning("Username already exists.")
+            return False
+        hashed_password = generate_password_hash(password)
+        users_collection.insert_one({"username": username, "password": hashed_password})
+        logger.info("User registration successful.")
+        return True
+
 # Initialize the bot instance
 logger.info("Initializing PoultryFarmBot instance.")
 bot = PoultryFarmBot(db)
@@ -207,8 +268,8 @@ if tokenizer.pad_token is None:
     tokenizer.add_special_tokens({'pad_token': '[PAD]'})
     model.resize_token_embeddings(len(tokenizer))
 
-# Llama 2 response generation
-def llama2_response(user_input):
+# Llama 3 response generation
+def llama3_response(user_input):
     """
     Generate a response using the Llama 2 model.
 
@@ -240,31 +301,43 @@ def llama2_response(user_input):
         return f"Error generating response: {str(e)}"
 
 # Main chatbot function
-def chatbot_response(image, text):
+def chatbot_response(image, text, username, password):
     """
     Handle user input and generate appropriate responses.
 
     Args:
         image (numpy.ndarray): Image input for disease detection.
         text (str): Text input for general queries.
+        username (str): Username for authentication.
+        password (str): Password for authentication.
 
     Returns:
         str: Response generated by the chatbot.
     """
+    user = bot.authenticate_user(username, password)
+    if not user:
+        return "Authentication failed. Please check your username and password."
+
+    user_id = user['_id']
+
     # If an image is provided, diagnose the disease
     if image is not None:
         logger.info("Image input detected. Proceeding with disease diagnosis.")
         diagnosis, name, status, recom = bot.diagnose_disease(image)
         if name and status and recom:
             logger.info("Diagnosis complete.")
+            bot.log_enquiry("image", "Image Enquiry", diagnosis, user_id)
             return diagnosis
         else:
             logger.warning("Diagnosis incomplete.")
+            bot.log_enquiry("image", "Image Enquiry", diagnosis, user_id)
             return diagnosis
     else:
         # Generate a response using Llama 2 for general text input
         logger.info("Text input detected. Generating response.")
-        return llama2_response(text)
+        response = llama3_response(text)
+        bot.log_enquiry("text", text, response, user_id)
+        return response
 
 # Gradio interface
 def build_gradio_interface():
@@ -294,6 +367,19 @@ def build_gradio_interface():
                     lines=3,
                     elem_id="user-input",
                 )
+                username = gr.Textbox(
+                    label="Username",
+                    placeholder="Enter your username",
+                    lines=1,
+                    elem_id="username-input",
+                )
+                password = gr.Textbox(
+                    label="Password",
+                    placeholder="Enter your password",
+                    type="password",
+                    lines=1,
+                    elem_id="password-input",
+                )
 
         output_box = gr.Textbox(
             label="Response",
@@ -311,7 +397,7 @@ def build_gradio_interface():
         # Connect the submit button to the chatbot response function
         submit_button.click(
             fn=chatbot_response,
-            inputs=[fecal_image, user_input],
+            inputs=[fecal_image, user_input, username, password],
             outputs=[output_box]
         )
     logger.info("Gradio interface built successfully.")