Update app.py

app.py

@@ -1,4 +1,255 @@
 import streamlit as st
+import uuid
+import sys
+import requests
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+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
+
+    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,
+    )
+
+    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:
+    """
+
+    from peft import *
+
+    # 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=print(generate_response(prompt))
+
+    # 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)