demospace / app.py
CogwiseAI's picture
Update app.py
322c3d1
raw
history blame
7.65 kB
import streamlit as st
import uuid
import sys
import requests
from peft import *
import bitsandbytes as bnb
import pandas as pd
import torch
import torch.nn as nn
import transformers
from datasets import load_dataset
from huggingface_hub import notebook_login
from peft import (
LoraConfig,
PeftConfig,
get_peft_model,
prepare_model_for_kbit_training,
)
from transformers import (
AutoConfig,
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,
)
USER_ICON = "images/user-icon.png"
AI_ICON = "images/ai-icon.png"
MAX_HISTORY_LENGTH = 5
if 'user_id' in st.session_state:
user_id = st.session_state['user_id']
else:
user_id = str(uuid.uuid4())
st.session_state['user_id'] = user_id
if 'chat_history' not in st.session_state:
st.session_state['chat_history'] = []
if "chats" not in st.session_state:
st.session_state.chats = [
{
'id': 0,
'question': '',
'answer': ''
}
]
if "questions" not in st.session_state:
st.session_state.questions = []
if "answers" not in st.session_state:
st.session_state.answers = []
if "input" not in st.session_state:
st.session_state.input = ""
st.markdown("""
<style>
.block-container {
padding-top: 32px;
padding-bottom: 32px;
padding-left: 0;
padding-right: 0;
}
.element-container img {
background-color: #000000;
}
.main-header {
font-size: 24px;
}
</style>
""", unsafe_allow_html=True)
def write_top_bar():
col1, col2, col3 = st.columns([1,10,2])
with col1:
st.image(AI_ICON, use_column_width='always')
with col2:
header = "Cogwise Intelligent Assistant"
st.write(f"<h3 class='main-header'>{header}</h3>", unsafe_allow_html=True)
with col3:
clear = st.button("Clear Chat")
return clear
clear = write_top_bar()
if clear:
st.session_state.questions = []
st.session_state.answers = []
st.session_state.input = ""
st.session_state["chat_history"] = []
def handle_input():
input = st.session_state.input
question_with_id = {
'question': input,
'id': len(st.session_state.questions)
}
st.session_state.questions.append(question_with_id)
chat_history = st.session_state["chat_history"]
if len(chat_history) == MAX_HISTORY_LENGTH:
chat_history = chat_history[:-1]
# api_url = "https://9pl792yjf9.execute-api.us-east-1.amazonaws.com/beta/chatcogwise"
# api_request_data = {"question": input, "session": user_id}
# api_response = requests.post(api_url, json=api_request_data)
# result = api_response.json()
# answer = result['answer']
# !pip install -Uqqq pip --progress-bar off
# !pip install -qqq bitsandbytes == 0.39.0
# !pip install -qqqtorch --2.0.1 --progress-bar off
# !pip install -qqq -U git + https://github.com/huggingface/transformers.git@e03a9cc --progress-bar off
# !pip install -qqq -U git + https://github.com/huggingface/peft.git@42a184f --progress-bar off
# !pip install -qqq -U git + https://github.com/huggingface/accelerate.git@c9fbb71 --progress-bar off
# !pip install -qqq datasets == 2.12.0 --progress-bar off
# !pip install -qqq loralib == 0.1.1 --progress-bar off
# !pip install einops
import os
# from pprint import pprint
# import json
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# notebook_login()
# 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
# """### Load dataset"""
from datasets import load_dataset
dataset_name = "nisaar/Lawyer_GPT_India"
# dataset_name = "patrick11434/TEST_LLM_DATASET"
dataset = load_dataset(dataset_name, split="train")
# """## Load adapters from the Hub
# You can also directly load adapters from the Hub using the commands below:
# """
# change peft_model_id
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
load_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
peft_model_id = "nisaar/falcon7b-Indian_Law_150Prompts"
config = PeftConfig.from_pretrained(peft_model_id)
model = AutoModelForCausalLM.from_pretrained(
config.base_model_name_or_path,
return_dict=True,
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
tokenizer.pad_token = tokenizer.eos_token
model = PeftModel.from_pretrained(model, peft_model_id)
"""## Inference
You can then directly use the trained model or the model that you have loaded from the 🤗 Hub for inference as you would do it usually in `transformers`.
"""
generation_config = model.generation_config
generation_config.max_new_tokens = 200
generation_config_temperature = 1
generation_config.top_p = 0.7
generation_config.num_return_sequences = 1
generation_config.pad_token_id = tokenizer.eos_token_id
generation_config_eod_token_id = tokenizer.eos_token_id
DEVICE = "cuda:0"
# Commented out IPython magic to ensure Python compatibility.
# %%time
# prompt = f"""
# <human>: Who appoints the Chief Justice of India?
# <assistant>:
# """.strip()
#
# encoding = tokenizer(prompt, return_tensors="pt").to(DEVICE)
# with torch.inference_mode():
# outputs = model.generate(
# input_ids=encoding.attention_mask,
# generation_config=generation_config,
# )
# print(tokenizer.decode(outputs[0],skip_special_tokens=True))
def generate_response(question: str) -> str:
prompt = f"""
<human>: {question}
<assistant>:
""".strip()
encoding = tokenizer(prompt, return_tensors="pt").to(DEVICE)
with torch.inference_mode():
outputs = model.generate(
input_ids=encoding.input_ids,
attention_mask=encoding.attention_mask,
generation_config=generation_config,
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
assistant_start = '<assistant>:'
response_start = response.find(assistant_start)
return response[response_start + len(assistant_start):].strip()
# prompt = "Debate the merits and demerits of introducing simultaneous elections in India?"
prompt=input
answer=generate_response(prompt)
print(answer)
# answer='Yes'
chat_history.append((input, answer))
st.session_state.answers.append({
'answer': answer,
'id': len(st.session_state.questions)
})
st.session_state.input = ""
def write_user_message(md):
col1, col2 = st.columns([1,12])
with col1:
st.image(USER_ICON, use_column_width='always')
with col2:
st.warning(md['question'])
def render_answer(answer):
col1, col2 = st.columns([1,12])
with col1:
st.image(AI_ICON, use_column_width='always')
with col2:
st.info(answer)
def write_chat_message(md, q):
chat = st.container()
with chat:
render_answer(md['answer'])
with st.container():
for (q, a) in zip(st.session_state.questions, st.session_state.answers):
write_user_message(q)
write_chat_message(a, q)
st.markdown('---')
input = st.text_input("You are talking to an AI, ask any question.", key="input", on_change=handle_input)