Update app.py

app.py

@@ -10,12 +10,10 @@ Original file is located at
 #!pip install gradio --quiet
 #!pip install -Uq transformers datasets timm accelerate evaluate
 
-import subprocess
-# subprocess.run('pip3 install datasets timm cv2 huggingface_hub torch pillow matplotlib' ,shell=True)
-
 import gradio as gr
-from huggingface_hub import hf_hub_download
 from safetensors.torch import load_model
+from timm import create_model
+from huggingface_hub import hf_hub_download
 from datasets import load_dataset
 import torch
 import torchvision.transforms as T
@@ -23,11 +21,18 @@ import cv2
 import matplotlib.pyplot as plt
 import numpy as np
 from PIL import Image
-from timm import create_model
+import os
+
+from langchain_community.document_loaders import TextLoader
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.text_splitter import CharacterTextSplitter
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
 from langchain_fireworks import ChatFireworks
-from langchain.schema import HumanMessage
-import os 
-
+from langchain_core.prompts import ChatPromptTemplate
+from transformers import AutoModelForImageClassification, AutoImageProcessor
+from langchain import HuggingFacePipeline
 
 safe_tensors = "model.safetensors" #hf_hub_download(repo_id="subhuatharva/swim-224-base-satellite-image-classification", filename="model.safetensors")
 
@@ -57,8 +62,59 @@ def one_hot_decoding(labels):
 title = "Satellite Image Classification for Landscape Analysis"
 description = """The bot was primarily trained to classify satellite images of the entire Amazon rainforest. You will need to upload satellite images and the bot will classify roads, forest, agriculure areas and the bot will return an analysis for the factors causing deforestation."""
 
+def ragChain():
+    """
+    function: creates a rag chain
+    output: rag chain
+    """
+    loader = TextLoader("document.txt")
+    docs = loader.load()
+
+    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
+    docs = text_splitter.split_documents(docs)
+
+    vectorstore = FAISS.load_local("faiss_index", embeddings = HuggingFaceEmbeddings(), allow_dangerous_deserialization = True)
+    retriever = vectorstore.as_retriever(search_type = "similarity", search_kwargs = {"k": 5})
+
+    api_key = os.getenv("FIREWORKS_API_KEY")
+    llm = ChatFireworks(model="accounts/fireworks/models/mixtral-8x7b-instruct", api_key = api_key)
+
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                """You are a knowledgeable landscape deforestation analyst. 
+                 """
+            ),
+            (
+                "human",
+                """First mention the detected labels only with short description.
+                Provide not more than 4 precautionary measures which are related to the detected labels that can be taken to control deforestation. 
+                Don't include conversational messages.
+                """,
+            ),
+
+            ("human", "{context}, {question}"),
+        ]
+    )
+
+    rag_chain = (
+    {
+    "context": retriever,
+    "question": RunnablePassthrough()
+    }
+    | prompt
+    | llm
+    | StrOutputParser()
+    )
+
+    return rag_chain
+
+
+
+
 def model_output(image):
-  #image = cv2.imread(image)
+
   PIL_image = Image.fromarray(image.astype('uint8'), 'RGB')
 
   img_size = (224,224)
@@ -76,15 +132,20 @@ def model_output(image):
   predictions = predictions.float().numpy().flatten()
   pred_labels = one_hot_decoding(predictions)
   output_text = " ".join(pred_labels)
-  query = f"summarize the classified satellite image labels {output_text} and summarize if or how these factors can cause deforestation. Don't write a conversational line."
-  
-  api_key = os.getenv("FIREWORKS_API_KEY")
-  llm = ChatFireworks(api_key=api_key, model="accounts/fireworks/models/mixtral-8x7b-instruct")
-  message = HumanMessage(query) 
 
-  return llm([message]).content
+  query = f"Detected labels in the provided satellite image are {output_text}. Give information on the labels."
+
+  return query
 
-app = gr.Interface(fn=model_output, inputs="image", outputs="text", title=title,
+def generate_response(rag_chain, query):
+    """
+    input: rag chain, query
+    function: generates response using llm and knowledge base
+    output: generated response by the llm
+    """
+    return rag_chain.invoke(f"{query}")
+    
+app = gr.Interface(fn=main, inputs="image", outputs="text", title=title,
     description=description, examples=[["sample_images/train_142.jpg"], ["sample_images/train_75.jpg"],["sample_images/train_32.jpg"], ["sample_images/train_59.jpg"], ["sample_images/train_67.jpg"], ["sample_images/train_92.jpg"], ["sample_images/random_satellite_image.png"]])
 app.launch(share=True)