# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="Y7W6CReNgcoI"
# ## Initialisation

# + id="_dx3DvFcESED"
# Check for running GPU
# !nvidia-smi

# + id="oawGovRjEkPF"
# Run once for initial installs
# #! pip install farm-haystack
# ! pip install git+https://github.com/deepset-ai/haystack.git
# ! pip install urllib3==1.25.4
# #! pip install torch==1.6.0+cu101 torchvision==0.6.1+cu101 -f https://download.pytorch.org/whl/torch_stable.html
# !pip install transformers[torch]
# !pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-1.7.0+cu101.html

# + id="JwGqL8YjCX_r"
# Run this if want to see graphic representation of pipeline 
# ! apt install libgraphviz-dev graphviz
# ! pip install pygraphviz

# + id="OINw-zqjFooP"
# Start Elasticsearch from source
# ! wget https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.6.2-linux-x86_64.tar.gz -q
# ! tar -xzf elasticsearch-7.6.2-linux-x86_64.tar.gz
# ! chown -R daemon:daemon elasticsearch-7.6.2
# ! sleep 30

# + colab={"base_uri": "https://localhost:8080/"} id="F1JXvtG1KFbu" outputId="5b7e08c8-483a-4116-8656-3659b47a2c18"
# Connect to Elasticsearch - Used for Elastic Retriever 
import os
from subprocess import Popen, PIPE, STDOUT
es_server = Popen(['elasticsearch-7.6.2/bin/elasticsearch'],
                   stdout=PIPE, stderr=STDOUT,
                   preexec_fn=lambda: os.setuid(1)
                  )
from haystack.document_store.elasticsearch import ElasticsearchDocumentStore
document_store_ES = ElasticsearchDocumentStore(similarity="dot_product", host="localhost", username="", password="")

# + id="Y845gghcCLeh"
from typing import List
import requests
import pandas as pd
from haystack import Document
from haystack.document_store.faiss import FAISSDocumentStore

# FAISS DocumentStore used for DPR and embedding retriever
document_store_FAISS = FAISSDocumentStore(
    faiss_index_factory_str="Flat",
    return_embedding=True
)

# + id="Sb566go43uNz"
import pprint
pp = pprint.PrettyPrinter(indent=2)

# + [markdown] id="dW60g6FZgjL0"
# ## Document Preprocessing

# + id="gT3fxP2UKHcS" colab={"base_uri": "https://localhost:8080/"} outputId="a895d482-67b0-4d0c-c448-b9c9d60373a2"
from haystack.reader.farm import FARMReader
import haystack

converter = haystack.file_converter.txt.TextConverter(
                    remove_numeric_tables=False,
                    valid_languages = ["en"])

as4 = converter.convert(file_path="/content/as4-winterBarley.txt")


processor = haystack.preprocessor.preprocessor.PreProcessor(
    clean_empty_lines=True,
    clean_whitespace=True,
    clean_header_footer=True,
    split_by="passage",
    split_length=1,
    split_respect_sentence_boundary=False,
    split_overlap=0
)

as4Docs = processor.process(as4)

for i in range(len(as4Docs)):
    as4Docs[i]["meta"]["table"] = False

document_store_ES.delete_all_documents()
document_store_ES.write_documents(as4Docs)
document_store_FAISS.delete_all_documents()
document_store_FAISS.write_documents(as4Docs)

#backagain = document_store.get_all_documents();

for i in range(0,len(as4Docs)):
    print(str(i) + ":", end = " ")
    print(as4Docs[i])

# + colab={"base_uri": "https://localhost:8080/"} id="pyRO7mk9pu7-" outputId="253b9b13-9947-479e-ef9b-14e0ea199eb5"
# Update table content and table description - csv for table content, txt for table description
# The files are under the /tables directory, the name does not matter, as long as the csv and the txt match each other.
# Need to rerun this section when new table is uploaded.

import csv
import os

data = []
docs= []

for file in [file for file in os.listdir(os.getcwd()+"/tables") if ".csv" in file]:
  with open(os.getcwd()+"/tables/"+file, mode='r') as infile:
    reader = csv.reader(infile)
    new_dict = {row[0]:row[1:] for row in reader}
    data.append(new_dict)
  infile.close()
  with open(os.getcwd()+"/tables/"+file[:-4]+".txt", mode='r') as infile:
    docs.append(infile.read())
  infile.close()

with open('table_text.txt', 'w') as outfile:
    for item in docs:
        outfile.write("%s\n\n" % item)
outfile.close()

# Construct FAISS DocumentStore for table content

tables = converter.convert(file_path="/content/table_text.txt")

tableDocs = processor.process(tables)
for i in range(len(tableDocs)):
  tableDocs[i]["meta"]["index"] = i
  tableDocs[i]["meta"]["table"] = True

document_store_ES.write_documents(tableDocs)
document_store_FAISS.write_documents(tableDocs)

# + colab={"base_uri": "https://localhost:8080/"} id="9ogJS95CIyng" outputId="002b3886-da5a-476d-ee17-deeb83d1fd2b"
# NLP based keyword extractor
# New category upload is not handled here
from KeyInfoExtractor import KeywordExtractor
key_word_extractor = KeywordExtractor()
key_word_extractor.read_key_info_file("categories.csv")
key_word_extractor.check_current_categories(True)
questions = {}
for ki in key_word_extractor.key_infos:
  questions[ki.name] = ki.question
pp.pprint(questions)

# + id="YYHZRJHfqai3" colab={"base_uri": "https://localhost:8080/"} outputId="8d46bae8-6dcc-4f8d-ba2c-a36f04b26eb9"
# Preliminary parser for questions - specially tackles the variety of the crops
# New json upload is not handled here
from parsing import Parser

parser = Parser()

# To parse question into general case: e.g. what to apply to Dunston -> what to apply to winter wheat
def question_parsing(text):
  parsed = parser.parse(text)
  parsed_string = ""
  for string in parsed["simple"]:
    parsed_string = parsed_string + string
  parsed_string = re.sub("[\(\[].*?[\)\]]", "", parsed_string)
  parsed_string = re.sub("  ", " ", parsed_string)
  return parsed_string

# To extract the keywords that the extractor cannot handle - e.g. variety
def generate_keywords(text):
  parsed = parser.parse(text)
  parsed_string = ""
  for string in parsed["simple"]:
    parsed_string = parsed_string + string
  keywords = {}
  while (parsed_string.find("(") != -1):
    start = parsed_string.find("(")
    end = parsed_string.find(")")
    bracket = parsed_string[start+1:end]
    if bracket.find(":") != -1:
      colon = bracket.find(":")
      keywords[bracket[1:colon].replace(" ", "")] = bracket[colon+1:].replace(" ", "")
    parsed_string = parsed_string[end+1:] 
  return keywords


# + id="zpl4rPhd2Wy-"
# Pressure score system when crops with different pressure score need to be treated differently.
# Need to rerun update pressure table if new pressure table is uploaded (need to include the origianl csv content)

import csv
import re

class PressureScoreGenerator:
  def __init__(self):
    self.pressure_table = {}

  def update_pressure_table(self,filename='pressure_score.csv'):
    with open('pressure_score.csv', mode='r') as infile:
        reader = csv.reader(infile)
        for row in reader:
          if row[0] in self.pressure_table:
            if row[1] in self.pressure_table[row[0]]:
              self.pressure_table[row[0]][row[1]][len(self.pressure_table[row[0]][row[1]])+1] = [row[2],int(row[3])]
            else:
              self.pressure_table[row[0]][row[1]] = {1:[row[2],int(row[3])]}
          else:
            self.pressure_table[row[0]]={row[1]:{1:[row[2],int(row[3])]}}
    infile.close()

  def calculate_pressure_score(self,crop):
    if crop in self.pressure_table:
      print("To give a better suggestion. We need to calculate the disease pressure of your current case. Please specify your choice by replying with the number accordingly.")
      score = 0
      # Would it be possible to change this into options by click?
      for key in self.pressure_table[crop].keys():
        print("Please choose the " + key + " which best describes your current case: ")
        for option in self.pressure_table[crop][key]:
          print(option, ": ",self.pressure_table[crop][key][option][0])
        choice = input()
        while not choice.isnumeric() or not (int(choice) >= 1 and int(choice) <= len(self.pressure_table[crop][key])):
          choice = input("Your input is not Valid. Please check and try again.")
        score = score + self.pressure_table[crop][key][int(choice)][1]
      if score <= 12:
        pressure_level = "low"
      elif score <= 15:
        pressure_level = "moderate"
      elif score <= 18:
        pressure_level = "high"
      else:
        pressure_level = "very high"
      print("Your current disease pressure score is " + str(score) + ".\nThis means that the disease pressure is " + pressure_level + ".")
      return pressure_level
    else:
      return -1

pressure_score_generator = PressureScoreGenerator()
pressure_score_generator.update_pressure_table()
#pp.pprint(pressure_score_generator.pressure_table)

# + [markdown] id="JrR-AGlxgqob"
# ## Building Individual Components

# + id="TFIjvlTwLSoN"
# Naive retriver based on tf * idf - Default BM25, can be cunstomised
from haystack.retriever.sparse import ElasticsearchRetriever
es_retriever = ElasticsearchRetriever(document_store=document_store_ES)

# + id="xMFr-plj4B0o"
# Alternative retriever - double BERT neural networks for question and doc embedding
from haystack.retriever.dense import DensePassageRetriever
dpr_retriever = DensePassageRetriever(
    document_store=document_store_FAISS,
    query_embedding_model="facebook/dpr-question_encoder-single-nq-base",
    passage_embedding_model="facebook/dpr-ctx_encoder-single-nq-base",
    use_gpu=True,
    embed_title=True,
)

# + colab={"base_uri": "https://localhost:8080/"} id="LVwNkA6TBRp1" outputId="b51c8da5-74f2-4a16-c6dd-fe5cd7672d1c"
document_store_FAISS.update_embeddings(retriever=dpr_retriever) #possible training of dpr model

# + colab={"base_uri": "https://localhost:8080/"} id="OuTEG8eN9UWm" outputId="8bcc0023-1e55-401e-fb5b-bd225907eff5"
# Alternative retriever - single BERT to embed both question and doc, may be better for similar documents (our case)
from haystack.retriever.dense import EmbeddingRetriever
embedding_retriever = EmbeddingRetriever(document_store=document_store_FAISS,
                               embedding_model="deepset/sentence_bert")


# + id="0sXdUWDpMzsw"
# Reader to further scan with Hugging Face models
# reader = TransformersReader(model_name_or_path="distilbert-base-uncased-distilled-squad", tokenizer="distilbert-base-uncased", use_gpu=-1)
# reader = FARMReader(model_name_or_path="deepset/bert-large-uncased-whole-word-masking-squad2", use_gpu=True)

# + id="MCN5zislt5j0"
# Decide whether the answer should be retrieved from the tables or the general texts
class QueryClassifier:
  outgoing_edges = 2

  def run(self, **kwargs):
    #print("Running Query Classifier")
    #print(len(kwargs["documents"]))
    #print(kwargs["documents"][0].meta)
    #for doc in kwargs["documents"]:
    #  print(doc.text)
    #print(kwargs["documents"][0].meta["table"] == "0")
    #docs = kwargs["documents"]
    #for doc in docs:
    #  print(doc.text)
    #  print(doc.meta)
    if (kwargs["documents"][0].meta["table"] == "0"):
      #print("It is a general question.")
      return (kwargs, "output_1")
    else:
      #print("It is a table-related question.")
      return (kwargs, "output_2")

query_classifier = QueryClassifier()

# + id="ursLVOh7-daI"
# Customised retriever - Process table content using TAPAS
from haystack.retriever.dense import BaseRetriever
from transformers import pipeline, AutoTokenizer, AutoModelForTableQuestionAnswering, TableQuestionAnsweringPipeline

class TableRetriever:
  outgoing_edges = 1

  def __init__(self):
    self.tokenizer = AutoTokenizer.from_pretrained("google/tapas-base-finetuned-wtq")
    self.model = AutoModelForTableQuestionAnswering.from_pretrained("google/tapas-base-finetuned-wtq")
    self.tableQA = TableQuestionAnsweringPipeline(model=self.model,tokenizer=self.tokenizer)

  def run(self,**kwargs):
    #print("Running Table Retriever")
    
    #print("Finished")
    return ({"result":self.tableQA(data[int(kwargs["documents"][0].meta["index"])], kwargs["query"])},"output_1")

table_retriever = TableRetriever()


# + id="gIkR1zR4DU__"
# Ask further questions if different context is detected.
# In the special case of pressure score, if it can be calculated, a different set of questions will be asked to decide the score.

class FurtherQuestionGenerator:
  outgoing_edges = 1

  def individualFiltersGenerator(self, text):
    text = text.lower()
    keywords = key_word_extractor.get_best_matches(text)
    #print(current_filters)
    #print(key_word_extractor.get_best_matches(text.lower(), give_list=True))
    
    # The extraction of pressure keywords work differently from others
    if text.find("pressure") != -1:
      keywords["pressure"] = []
      if "low" in text[text.find("pressure")-20:text.find("pressure")+30]:
        keywords["pressure"].append("low")
      if "moderate" in text[text.find("pressure")-20:text.find("pressure")+30]:
        keywords["pressure"].append("moderate")
      if "high" in text[text.find("pressure")-20:text.find("pressure")+30]:
        keywords["pressure"].append("high")
      if "very high" in text[text.lower().find("pressure")-20:text.lower().find("pressure")+30] or "extreme" in text[text.find("pressure")-20:text.find("pressure")+30]:
        keywords["pressure"].append("very high")
    
    # To include special keywords such as variety
    parsed_keywords = generate_keywords(text)
    if parsed_keywords != {}:
      for key,value in parsed_keywords.items():
        if key not in keywords:
          keywords[key] = [value]
        else:
          keywords[key].append(value)
    return keywords

  def topDocsFilterGenerator(self, docs):
    return [self.individualFiltersGenerator(doc.text) for doc in docs]

  def filters_difference(self, filters_list, specified = []):
    current_filters = {}
    for filters in filters_list:
      for category, filters in filters.items():
        if ((category not in specified) and (category in current_filters) and (filters != current_filters[category])):
          return category
        elif ((category not in specified) and (category not in current_filters)):
          current_filters[category] = filters
    return None

  def furtherQuestions(self, docs, specified=[], original_filters={}):
    filters_list = self.topDocsFilterGenerator(docs)
    #print(len(docs),len(filters_list))

    # Filter the retrieved docs before asking questions, eliminate docs with different keywords but keep those does not mention the keywords.
    for keyword,new_key in original_filters.items():
      temp_d = [doc for i,doc in enumerate(docs) if not ((keyword in filters_list[i].keys()) and (new_key[0].lower() not in filters_list[i][keyword]))]
      temp_f = [filters for i,filters in enumerate(filters_list) if not ((keyword in filters_list[i].keys()) and (new_key[0].lower() not in filters_list[i][keyword]))]
    docs = temp_d
    filters_list = temp_f
    #print(len(docs),len(filters_list))

    # Rank the docs based on closeness to keywords in question
    match = [0 for doc in docs]
    for keyword,new_key in original_filters.items():
      match = [match[i] + 1 if ((keyword in filters_list[i].keys()) and (new_key[0].lower() in filters_list[i][keyword])) else match[i] for i in range(len(filters_list))]
      match = [match[i] + 1 if ((keyword in filters_list[i].keys()) and (filters_list[i][keyword] == [new_key[0].lower()])) else match[i] for i in range(len(filters_list))]
    
    # Use only questions from the top 3 results. Otherwise too many questions asked.
    docs = [x for _,x in sorted(zip(match,docs), key=lambda pair: pair[0], reverse=True)]
    filters_list = [x for _,x in sorted(zip(match,filters_list), key=lambda pair: pair[0], reverse=True)]
    keyword = self.filters_difference(filters_list[:3],specified)
    contained_pressure = False
    while keyword is not None:
      if keyword != "pressure":
        new_key = input(questions[keyword] + "(If not, please reply No) ")
        if (new_key.lower() != "no"):
          match = [match[i] + 1 if ((keyword in filters_list[i].keys()) and (new_key.lower() in filters_list[i][keyword])) else match[i] for i in range(len(filters_list))]
          match = [match[i] + 1 if ((keyword in filters_list[i].keys()) and (filters_list[i][keyword] == [new_key.lower()])) else match[i] for i in range(len(filters_list))]
        specified.append(keyword)
        original_filters[keyword] = [new_key]
      else:
        contained_pressure = True
        specified.append(keyword)
        original_filters[keyword] = []
      keyword = self.filters_difference(filters_list[:3],specified)
    
    # Treat the special set of pressure score questions
    if contained_pressure and "Crop" in original_filters:
        pressure_level = pressure_score_generator.calculate_pressure_score(original_filters["Crop"][0])
        match = [match[i] + 1 if (("pressure" in filters_list[i].keys()) and (pressure_level in filters_list[i]["pressure"])) else match[i] for i in range(len(filters_list))]
        match = [match[i] + 1 if (("pressure" in filters_list[i].keys()) and (filters_list[i]["pressure"] == [pressure_level])) else match[i] for i in range(len(filters_list))]
    sorted_docs = sorted(zip(match,docs), key=lambda pair: pair[0], reverse=True)
    #for doc in sorted_docs:
      #print(doc[1].text)
      #print(doc[1].meta)
    return sorted_docs[:3]

  def run(self, **kwargs):
    #print("Running Question Generator")
    original_filters = self.individualFiltersGenerator(kwargs["query"])
    specified = list(original_filters.keys())
    return ({"result": self.furtherQuestions(kwargs["documents"],specified,original_filters)},"output_1")

question_generator = FurtherQuestionGenerator()


# + id="7wugdBgtxoaH"
class Result:
  outgoing_edges = 1

  def run(self,**kwargs):
    #print("Running Result")
    #print(type(kwargs["result"]))
    if type(kwargs["result"]) is list:
      docs = [doc[1] for doc in kwargs["result"]]
      return (docs,"output_1")
    else:
      return (kwargs["result"]["answer"],"output_1")

result = Result()

# + [markdown] id="OHatnDA3g2ON"
# ## Assembling into Pipeline

# + colab={"base_uri": "https://localhost:8080/"} id="HEH2V1xi3Hw5" outputId="83999776-c911-48d7-c2ed-07866a28fd86"
# Current approach
from haystack import Pipeline
from haystack.pipeline import JoinDocuments

# Building new pipeline with multiple retrievers
pipeline = Pipeline()
pipeline.add_node(component=es_retriever, name="ESRetriever", inputs=["Query"])
pipeline.add_node(component=dpr_retriever, name="DPRRetriever", inputs=["Query"])
pipeline.add_node(component=embedding_retriever, name="EmbeddingRetriever", inputs=["Query"])
pipeline.add_node(component=JoinDocuments(join_mode="merge"), name="JoinResults", inputs=["DPRRetriever","EmbeddingRetriever","ESRetriever"])
pipeline.add_node(component=query_classifier, name="QueryClassifier", inputs=["JoinResults"])
pipeline.add_node(component=question_generator, name="QnGenerator", inputs=["QueryClassifier.output_1"])
pipeline.add_node(component=table_retriever, name="TableRetriever", inputs=["QueryClassifier.output_2"])
pipeline.add_node(component=result, name="Result", inputs=["QnGenerator", "TableRetriever"])
pipeline.draw(path="custom_pipe.png")

# Question input goes here
question = input("What would you like to ask about?")
response = pipeline.run(query=question_parsing(question), top_k_retriever=20)

# Final answer is here 
#final_result = response.text if type(response) is Document else response["answer"]
if type(response) is list:
  print(response[0].text)
else:
  print(response)

StarcoderdataJupyterScriptsDedupFiltered