# ---
# 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
# ---

# + id="NgfUN_Dj4Vc4" executionInfo={"status": "ok", "timestamp": 1601198903486, "user_tz": -480, "elapsed": 12302, "user": {"displayName": "\u5982\u5b50", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi3ItGjzEGzUOlXTUHjOgeuVA5TICdNcY-Q1TGicA=s64", "userId": "01997730851420384589"}}
# !wget https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.9.0-linux-x86_64.tar.gz -q
# !tar -xzf elasticsearch-7.9.0-linux-x86_64.tar.gz
# !chown -R daemon:daemon elasticsearch-7.9.0

# + id="C672zBF84r1s" executionInfo={"status": "ok", "timestamp": 1601198915081, "user_tz": -480, "elapsed": 21380, "user": {"displayName": "\u5982\u5b50", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi3ItGjzEGzUOlXTUHjOgeuVA5TICdNcY-Q1TGicA=s64", "userId": "01997730851420384589"}} outputId="b79efc70-ce90-46fe-db34-d6491823593a" colab={"base_uri": "https://localhost:8080/", "height": 121}
# !pip install elasticsearch -q
# !pip install transformers -q

# + id="Ag4v3XUM4tys" executionInfo={"status": "ok", "timestamp": 1601198920952, "user_tz": -480, "elapsed": 26136, "user": {"displayName": "\u5982\u5b50", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi3ItGjzEGzUOlXTUHjOgeuVA5TICdNcY-Q1TGicA=s64", "userId": "01997730851420384589"}}
from subprocess import Popen, PIPE, STDOUT
from elasticsearch import Elasticsearch
from elasticsearch import helpers
from transformers import BertTokenizer

import os
import time
import pprint
import numpy as np
import tensorflow as tf
import tensorflow_hub as hub

# + id="duAm-4xe4vnN" executionInfo={"status": "ok", "timestamp": 1601198920962, "user_tz": -480, "elapsed": 24321, "user": {"displayName": "\u5982\u5b50", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi3ItGjzEGzUOlXTUHjOgeuVA5TICdNcY-Q1TGicA=s64", "userId": "01997730851420384589"}}
es_server = Popen(
  ['elasticsearch-7.9.0/bin/elasticsearch'], 
  stdout = PIPE, stderr = STDOUT,
  preexec_fn = lambda: os.setuid(1))


# + id="OuX2cpdR4xdt" executionInfo={"status": "ok", "timestamp": 1601198953943, "user_tz": -480, "elapsed": 1586, "user": {"displayName": "\u5982\u5b50", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi3ItGjzEGzUOlXTUHjOgeuVA5TICdNcY-Q1TGicA=s64", "userId": "01997730851420384589"}} outputId="15ccf875-15a4-4cde-df9c-dcb05560e011" colab={"base_uri": "https://localhost:8080/", "height": 312}
# !curl -X GET "localhost:9200/"

# + id="x7z9s-eA5XS4" executionInfo={"status": "ok", "timestamp": 1601198958188, "user_tz": -480, "elapsed": 1800, "user": {"displayName": "\u5982\u5b50", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi3ItGjzEGzUOlXTUHjOgeuVA5TICdNcY-Q1TGicA=s64", "userId": "01997730851420384589"}}
def gen_data():
  with open('/content/gdrive/My Drive/finch/es/free_chat/data/basic.txt') as f:
    for line in f:
      line = line.rstrip()
      q, a = line.split('<SEP>')
      bert_inp = ['[CLS]'] + list(q) + ['[SEP]']
      bert_inp = tokenizer.convert_tokens_to_ids(bert_inp)
      bert_seg = [0] * len(bert_inp)
      bert_mask = [1] * len(bert_inp)
      _, res = encoder([np.asarray([bert_inp], dtype=np.int32),
                        np.asarray([bert_mask], dtype=np.int32),
                        np.asarray([bert_seg], dtype=np.int32)])
      yield {
        '_index': 'chatbot',
        'question': q,
        'answer': a,
        'question_embedding': tf.reduce_mean(res, axis=1)[0].numpy(),}


# + id="Rja_ufvA8bVb" executionInfo={"status": "ok", "timestamp": 1601198980171, "user_tz": -480, "elapsed": 20904, "user": {"displayName": "\u5982\u5b50", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gi3ItGjzEGzUOlXTUHjOgeuVA5TICdNcY-Q1TGicA=s64", "userId": "01997730851420384589"}} outputId="102be757-03bb-4dc3-ab35-9c84825ca8a6" colab={"base_uri": "https://localhost:8080/", "height": 34}
from google.colab import drive
drive.mount('/content/gdrive')

# + id="HjUrl2cwV6id"
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese',
                                           lowercase = True,
                                           add_special_tokens = True)

encoder = hub.KerasLayer("https://tfhub.dev/tensorflow/bert_zh_L-12_H-768_A-12/2",
                         trainable=False)

# + id="5f85t8H641cG" executionInfo={"status": "ok", "timestamp": 1601199051617, "user_tz": -480, "elapsed": 53330, "user": {"displayName": "\u5982\u5b50", "photoUrl": "https://lh3.googleusercontent.com/a-/AO<KEY>Y-Q1TGicA=s64", "userId": "01997730851420384589"}} outputId="88952e4f-033e-4752-bab6-6ce59af4730b" colab={"base_uri": "https://localhost:8080/", "height": 52}
es = Elasticsearch()
print(es.ping())

es.indices.create(index='chatbot')

mapping = {
  'properties': {
    'question': {
      'type': 'text',
    },
    'question_embedding': {
      'type': 'dense_vector',
      'dims': 768,
    },
  }
}
es.indices.put_mapping(body=mapping, index='chatbot')

helpers.bulk(es, gen_data())

# + id="stFimp4m6FXj" outputId="e66fc56e-fe84-47ac-975f-462cd3f02fc1" colab={"base_uri": "https://localhost:8080/", "height": 1000}
while True:
  text_inp = input('Input:')
  t0 = time.time()
  bert_inp = ['[CLS]'] + list(text_inp) + ['[SEP]']
  bert_inp = tokenizer.convert_tokens_to_ids(bert_inp)
  bert_seg = [0] * len(bert_inp)
  bert_mask = [1] * len(bert_inp)
  _, res = encoder([np.asarray([bert_inp], dtype=np.int32),
                    np.asarray([bert_mask], dtype=np.int32),
                    np.asarray([bert_seg], dtype=np.int32)])
  query_vector = tf.reduce_mean(res, axis=1)[0].numpy()
  script_query = {
  'script_score': {
    'query': {'match_all': {}},
    'script': {
      'source': "cosineSimilarity(params.query_vector, doc['question_embedding']) + 1.0",
      'params': {'query_vector': query_vector},
      }
    }
  }
  dsl = {
    'query': script_query,
    '_source': {'excludes': ['question_embedding']},
  }
  hits = es.search(index='chatbot', body=dsl)['hits']['hits']
  print('Output:', hits[0]['_source']['answer'])
  print('%.2f sec' % (time.time() - t0))
  print()

StarcoderdataJupyterScriptsDedupFiltered