# -*- coding: utf-8 -*-
"""Assistant.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1tbapNWoymz_Oq_M96lTAeUDZrvKcB2ji
"""

#Define Web Scraping function
def web_scrape(url, tag, id):
  headers = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36",
  }
  page = requests.get(url, headers=headers)
  soup = BeautifulSoup(page.content, 'html.parser')
  data = soup.find_all(tag, class_ = id)
  res = []
  i = 0
  while i < 6:
    res.append(data[i].get_text())
    i = i+1
  return res

#Define Google result scraping function
def google_scrape(query):
  url = "https://www.google.com/search?q=" + str(query)
  res = web_scrape(url, "div", "VwiC3b yXK7lf MUxGbd yDYNvb lyLwlc lEBKkf")
  return res

#Define Bratanica web scraping function
def bratanica_scrape(query):
  headers = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36",
  }
  url = "https://www.britannica.com"
  page = requests.get(url + "/search?query=" + query, headers=headers)
  soup = BeautifulSoup(page.content, 'html.parser')
  link = soup.find_all("a", class_ ="font-weight-bold font-18")
  page = requests.get(url+link[0]["href"], headers=headers)
  soup = BeautifulSoup(page.content, 'html.parser')
  data = soup.find_all("p")
  res = []
  for i in data:
    res.append(i.get_text())
  return res

import requests
from bs4 import BeautifulSoup

import torch
from transformers import LlamaForCausalLM, LlamaTokenizer

# Hugging Face model_path
model_path = 'psmathur/orca_mini_3b'
tokenizer = LlamaTokenizer.from_pretrained(model_path)
model = LlamaForCausalLM.from_pretrained(
    model_path, torch_dtype=torch.float16, device_map='auto', load_in_4bit = True
)

#generate response function
def generate_text(system, instruction, gen_len, input=None):

    if input:
        prompt = f"### System:\n{system}\n\n### User:\n{instruction}\n\n### Context:\n{input}\n\n### Response:\n"
    else:
        prompt = f"### System:\n{system}\n\n### User:\n{instruction}\n\n### Response:\n"

    tokens = tokenizer.encode(prompt)
    tokens = torch.LongTensor(tokens).unsqueeze(0)
    tokens = tokens.to('cuda')

    instance = {'input_ids': tokens,'top_p': 1.5, 'temperature':0.3, 'generate_len': gen_len, 'top_k': 200}

    length = len(tokens[0])
    with torch.no_grad():
        rest = model.generate(
            input_ids=tokens,
            max_length=length+instance['generate_len'],
            use_cache=True,
            do_sample=True,
            top_p=instance['top_p'],
            temperature=instance['temperature'],
            top_k=instance['top_k']
        )
    output = rest[0][length:]
    string = tokenizer.decode(output, skip_special_tokens=True)
    return f'{string}'

def context(query):
  headers = {
      'Authorization': 'Bearer FGSCPX4IMD3MGEQHWOORJDR73RHT7LFZ',
  }

  params = {
      'v': '20230820',
      'q': query,
  }

  response = requests.get('https://api.wit.ai/message', params=params, headers=headers).json()

  return response["intents"][0]["name"]

convo = []
while True:
  q = input("Human:")
  convo.append(convo)

  if (context(q)) == "Does_not_require_context":
    system = "You are an AI chatbot that aims to be as helpful and funny as possible. You will be given the past conversation history to know what you were talking about before"
    instruction = q
    res = generate_text(system, instruction, 512, convo)
    convo.append("Bot:" + res)
    print(res)
  elif (context(q)) == "Requires_additional_context":
    system = "You are an AI chatbot that aims to be as helpful as possible. To do so, using the context, answer the question that the user asks you."
    instruction = q
    clue = google_scrape(q)
    res = generate_text(system, instruction, 512, clue[0])
    convo.append("Bot:" + res)
    print(res)