"Update space"

Signed-off-by: Emmanuel Frimpong Asante <[email protected]>

app.py

@@ -19,7 +19,7 @@ auth_model = load_model('models/auth_model.h5', compile=True)
 # Check for GPU availability and load the model accordingly
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load the tokenizer and model, ensuring they run on the correct device
+# Load the tokenizer and LLaMA model, ensuring they run on the correct device
 llama_tokenizer = AutoTokenizer.from_pretrained('meta-llama/Meta-Llama-3-8B-Instruct')
 llama_model = AutoModelForCausalLM.from_pretrained(
     'meta-llama/Meta-Llama-3-8B-Instruct',
@@ -48,46 +48,49 @@ def predict(image):
         status = result.get(indx)
         recom = recommend.get(indx)
 
-        return f"The chicken is in a {status} condition, diagnosed with {name}. The recommended medication is {recom}."
+        diagnosis = f"The chicken is in a {status} condition, diagnosed with {name}. The recommended medication is {recom}."
+        return diagnosis, name, status, recom
     else:  # If the image is not recognized as a chicken disease image
         return (
             "The uploaded image is not recognized as a chicken or does not appear to be related to any known chicken diseases. "
             "Please ensure the image is clear and shows a chicken or its symptoms to receive a proper diagnosis."
-        )
+        ), None, None, None
 
 
-def chat_response(user_input):
-    inputs = llama_tokenizer(user_input, return_tensors='pt').to(device)  # Ensure tensors are on the right device
-    outputs = llama_model.generate(inputs['input_ids'], max_length=500, do_sample=True)
-    response = llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
-
-
-def combined_interface(image, text):
+def generate_combined_response(image, text):
     if image is not None:
-        return predict(image)
+        diagnosis, name, status, recom = predict(image)
+
+        if name and status and recom:  # If the disease is recognized
+            # Generate a response using LLaMA based on the diagnosis
+            context = f"The chicken is in a {status} condition, diagnosed with {name}. The recommended medication is {recom}. "
+            if text:
+                context += f"Additionally, the user asked: '{text}'"
+            inputs = llama_tokenizer(context, return_tensors='pt').to(device)
+            outputs = llama_model.generate(inputs['input_ids'], max_length=500, do_sample=True)
+            advice = llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
+            return diagnosis + "\n\nAdditional Advice: " + advice
+        else:
+            return diagnosis
     elif text:
+        # Only text input provided, generate response using LLaMA
         return chat_response(text)
     else:
         return "Please provide an image or ask a question."
 
 
-example_image_path = 'disease.jpg'  # Replace with actual file path
-
-# Check if the example image exists and is a file
-if not os.path.isfile(example_image_path):
-    print(f"Warning: Example image file not found: {example_image_path}")
-    example_image_path = None  # Disable the example
+def chat_response(user_input):
+    inputs = llama_tokenizer(user_input, return_tensors='pt').to(device)
+    outputs = llama_model.generate(inputs['input_ids'], max_length=500, do_sample=True)
+    response = llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
 
-# Only include the example if the path is valid
-examples = [[example_image_path, ''], ['', 'What should I do if my chicken is sick?']] if example_image_path else None
 
-# Create Gradio interface
+# Gradio Interface
 interface = gr.Interface(
-    fn=combined_interface,
+    fn=generate_combined_response,
     inputs=[gr.Image(label='Upload Image'), gr.Textbox(label='Ask a question')],
-    outputs=gr.Textbox(label="Response"),
-    # examples=examples  # Use examples only if valid
+    outputs=gr.Textbox(label="Response")
 )
 
 # Launch the interface