Update app.py

app.py

@@ -1,29 +1,33 @@
 import gradio as gr
 import copy
+import os
 import time
 import llama_cpp
 from llama_cpp import Llama
 from huggingface_hub import hf_hub_download
-import os
-from gradio.components import Image, Text
-
+from fastai.vision.all import *
 
+# Load the LLM model
 llm = Llama(
     model_path=hf_hub_download(
         repo_id=os.environ.get("REPO_ID", "TheBloke/Llama-2-7B-Chat-GGML"),
         filename=os.environ.get("MODEL_FILE", "llama-2-7b-chat.ggmlv3.q5_0.bin"),
     ),
     n_ctx=2048,
-    n_gpu_layers=50, # change n_gpu_layers if you have more or less VRAM 
+    n_gpu_layers=50,  # change n_gpu_layers if you have more or less VRAM
 )
 
 history = []
 
 system_message = """
-You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
-If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+You are a BIRD EXPERT
 """
+# The rest of the system message
+
+# Load the Vision Model
+learn = load_learner('export.pkl')
 
+labels = learn.dls.vocab
 
 def generate_text(message, history):
     temp = ""
@@ -37,13 +41,13 @@ def generate_text(message, history):
         input_prompt,
         temperature=0.15,
         top_p=0.1,
-        top_k=40, 
+        top_k=40,
         repeat_penalty=1.1,
         max_tokens=1024,
         stop=[
-            "<|prompter|>",
-            "<|endoftext|>",
-            "<|endoftext|> \n",
+            "",
+            "",
+            " \n",
             "ASSISTANT:",
             "USER:",
             "SYSTEM:",
@@ -57,28 +61,35 @@ def generate_text(message, history):
 
     history = ["init", input_prompt]
 
-
 def predict(img):
     try:
         img = PILImage.create(img)
     except:
         return {"bird": "Unknown"}
-    pred,pred_idx,probs = learn.predict(img)
-    return {labels[i]: float(probs[i]) for i in range(len(labels))}
+    pred, pred_idx, probs = learn.predict(img)
+    return {"bird": labels[pred_idx], "probs": {labels[i]: float(probs[i]) for i in range(len(labels))}}
 
-title = "Bird Detector"
-description = "Bird Detector."
-examples = ['BIRD.png']
-interpretation='default'
-enable_queue=True
+title = "Bird Detector with LLM"
+description = "Detect birds and get LLM responses."
+examples = [{"img": "BIRD.jpg", "message": "Tell me about the bird."}]
+interpretation = 'default'
+enable_queue = True
 
 def combined(img, message):
     prediction = predict(img)
-    response = generate_text(message, history)
-    if "I have detected" in response:
-        response = response.replace("I have detected", f"I have detected {prediction['bird']} in the image.")
-
+    response = list(generate_text(f"I have detected {prediction['bird']} in the image. {message}", history))
     return response
 
-
-gr.Interface(fn=combined,inputs=Image(shape=(512, 512)),outputs=Text(),title=title,description=description,examples=examples,interpretation=interpretation,enable_queue=enable_queue).launch()
+gr.Interface(
+    fn=combined,
+    inputs=[
+        gr.inputs.Image(),
+        gr.inputs.Textbox(label="Message to LLM")
+    ],
+    outputs=gr.outputs.Textbox(),
+    title=title,
+    description=description,
+    examples=examples,
+    interpretation=interpretation,
+    enable_queue=enable_queue,
+).launch()