Update app.py

app.py

@@ -1,5 +1,258 @@
 import streamlit as st
+import uuid
+import sys
+import requests
+from peft import *
+import bitsandbytes as bnb
 import pandas as pd
-import pickle
+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,
+)
 
-model = pickle.load(open('model_saved.pkl', 'rb'))
+
+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)