Zwea Htet commited on
Commit
0809507
1 Parent(s): 0e8a2bc

added codes

Browse files
Files changed (5) hide show
  1. .gitignore +5 -0
  2. app.py +55 -2
  3. models/bloom.py +63 -0
  4. requirements.txt +10 -0
  5. utils/customLLM.py +36 -0
.gitignore ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ venv
2
+ data/__pycache__
3
+ models/__pycache__
4
+ .env
5
+ __pycache__
app.py CHANGED
@@ -1,4 +1,57 @@
 
 
 
1
  import streamlit as st
2
 
3
- x = st.slider('Select a value')
4
- st.write(x, 'squared is', x * x)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ import requests
4
  import streamlit as st
5
 
6
+ from models import bloom
7
+
8
+ st.title("Welcome to RegBotBeta")
9
+ st.header("Powered by `LlamaIndex🦙` and `OpenAI API`")
10
+
11
+ def validate(token: str):
12
+ api_endpoint = "https://api.openai.com/v1/chat/completions"
13
+ api_key = token
14
+
15
+ headers = {
16
+ "Content-Type" : "application/json",
17
+ "Authorization": f"Bearer {api_key}"
18
+ }
19
+
20
+ messages = [
21
+ {"role": "user", "content": "Say this is a test!"}
22
+ ]
23
+
24
+ data = {
25
+ "model": "gpt-3.5-turbo",
26
+ "messages": messages
27
+ }
28
+
29
+ response = requests.post(api_endpoint, json=data, headers=headers)
30
+ return response
31
+
32
+ def create_index():
33
+ index = bloom.initialize_index("")
34
+
35
+
36
+ def get_response(vector_index, query_str):
37
+ query_engine = vector_index.as_query_engine()
38
+ response = query_engine.query(query_str)
39
+ return response
40
+
41
+ api_key = st.text_input("Enter your OpenAI API key here:", type="password")
42
+ if api_key:
43
+
44
+
45
+ st.write("---")
46
+ input_text = st.text_area("Ask your question")
47
+
48
+ if input_text is not None:
49
+ if st.button("Ask"):
50
+ st.info("Your query: \n" + input_text)
51
+ with st.spinner("Processing your query..."):
52
+ response = get_response(index, input_text)
53
+ print(response)
54
+
55
+ st.success(response)
56
+
57
+ st.write("---")
models/bloom.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from json import dumps, loads
3
+
4
+ import numpy as np
5
+ import pandas as pd
6
+ from dotenv import load_dotenv
7
+ from llama_index import (Document, GPTVectorStoreIndex, LLMPredictor,
8
+ PromptHelper, ServiceContext, StorageContext,
9
+ load_index_from_storage)
10
+ from transformers import AutoModelForCausalLM, AutoTokenizer
11
+
12
+ from utils.customLLM import CustomLLM
13
+
14
+ load_dotenv()
15
+
16
+ # get model
17
+ model_name = "bigscience/bloom-560m"
18
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
19
+ model = AutoModelForCausalLM.from_pretrained(model_name, config='T5Config')
20
+
21
+ # define prompt helper
22
+ # set maximum input size
23
+ max_input_size = 2048
24
+ # set number of output tokens
25
+ num_output = 525
26
+ # set maximum chunk overlap
27
+ max_chunk_overlap = 20
28
+ prompt_helper = PromptHelper(max_input_size, num_output, max_chunk_overlap)
29
+
30
+ # define llm
31
+ llm_predictor = LLMPredictor(llm=CustomLLM(model, tokenizer))
32
+ service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
33
+
34
+ def prepare_data(file_path:str):
35
+ df = pd.read_json(file_path)
36
+ df = df.replace(to_replace="", value=np.nan).dropna(axis=0) # remove null values
37
+
38
+ parsed = loads(df.to_json(orient="records"))
39
+
40
+ documents = []
41
+ for item in parsed:
42
+ document = Document(item['paragraphText'],
43
+ item['_id']['$oid'],
44
+ extra_info={"chapter": item['chapter'],
45
+ "article": item['article'],
46
+ "title": item['title']})
47
+ documents.append(document)
48
+
49
+ return documents
50
+
51
+ def initialize_index(index_name):
52
+ file_path = f"./vectorStores/{index_name}"
53
+ if os.path.exists(file_path):
54
+ # rebuild storage context
55
+ storage_context = StorageContext.from_defaults(persist_dir=file_path)
56
+ # load index
57
+ index = load_index_from_storage(storage_context)
58
+ return GPTVectorStoreIndex.load_from_disk(file_path)
59
+ else:
60
+ documents = prepare_data(r"./assets/regItems.json")
61
+ index = GPTVectorStoreIndex.from_documents(documents, service_context=service_context)
62
+ index.storage_context.persist(file_path)
63
+ return index
requirements.txt ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ llama_index
2
+ torch
3
+ transformers
4
+ panda
5
+ numpy
6
+ langchain
7
+ openai
8
+ faiss-cpu
9
+ python-dotenv
10
+ streamlit
utils/customLLM.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, List, Mapping, Optional
2
+
3
+ from langchain.llms.base import LLM
4
+ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
5
+
6
+
7
+ class CustomLLM(LLM):
8
+
9
+ # Create the pipeline for question answering
10
+ def __init__(self, model: AutoModelForCausalLM, tokenizer: AutoTokenizer):
11
+ self.pipeline = pipeline(
12
+ model=model,
13
+ tokenizer=tokenizer,
14
+ task="text-generation",
15
+ # device=0, # GPU device number
16
+ # max_length=512,
17
+ do_sample=True,
18
+ top_p=0.95,
19
+ top_k=50,
20
+ temperature=0.7
21
+ )
22
+
23
+ def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
24
+ prompt_length = len(prompt)
25
+ response = self.pipeline(prompt, max_new_tokens=525)[0]["generated_text"]
26
+
27
+ # only return newly generated tokens
28
+ return response[prompt_length:]
29
+
30
+ @property
31
+ def _identifying_params(self) -> Mapping[str, Any]:
32
+ return {"name_of_model": self.model_name}
33
+
34
+ @property
35
+ def _llm_type(self) -> str:
36
+ return "custom"