app18

app.py

@@ -31,7 +31,7 @@ from pydantic import BaseModel
 import shutil
 
 # Cell 1: Image Classification Model
-image_pipeline = pipeline(task="image-classification", model="microsoft/resnet-50")
+image_pipeline = pipeline(task="image-classification", model="rocioadlc/TrashNet_ResNet152V2")
 
 def predict_image(input_img):
     predictions = image_pipeline(input_img)
@@ -68,10 +68,11 @@ vectordb = Chroma.from_documents(
     persist_directory=persist_directory
 )
 # define retriever
-retriever = vectordb.as_retriever(search_kwargs={"k": 3},search_type="mmr")
+retriever = vectordb.as_retriever(search_type="mmr")
 prompt_template = """
 Your name is AngryGreta and you are a recycling chatbot with the objective to anwer questions from user in English or Spanish /
 Use the following pieces of context to answer the question if the question is related with recycling /
+No more than two chunks of context /
 Answer in the same language of the question /
 Always say "thanks for asking!" at the end of the answer /
 If the context is not relevant, please answer the question by using your own knowledge about the topic.
@@ -101,7 +102,7 @@ llm = HuggingFaceHub(
     },
 )
 
-memory = ConversationBufferMemory(llm=llm, memory_key="chat_history", input_key='question', output_key='answer')
+memory = ConversationBufferMemory(llm=llm, memory_key="chat_history", input_key='question', output_key='answer', return_messages=True)
 
 qa_chain = ConversationalRetrievalChain.from_llm(
     llm = llm,
@@ -110,17 +111,14 @@ qa_chain = ConversationalRetrievalChain.from_llm(
     verbose = True,
     combine_docs_chain_kwargs={'prompt': qa_prompt},
     get_chat_history = lambda h : h,
-    rephrase_question = False
+    rephrase_question = False,
+    output_key = 'answer'
 )
 
-def chat_interface(question):
+def chat_interface(question,history):
 
     result = qa_chain.invoke({"question": question})
-    
-    # Extract only the answer from the result
-    answer = result.get('answer', None)
-    
-    return answer
+    return result['answer']
 
 chatbot_gradio_app = gr.ChatInterface(
     fn=chat_interface,