Zwea Htet commited on
Commit
3b7cf58
1 Parent(s): 8fe964e

fixed some bugs

Browse files
Files changed (3) hide show
  1. app.py +14 -38
  2. models/bloom.py +15 -15
  3. utils/customLLM.py +17 -5
app.py CHANGED
@@ -2,62 +2,38 @@ 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("bloomLlama")
34
- return index
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
 
42
  api_key = st.text_input("Enter your OpenAI API key here:", type="password")
43
  if api_key:
44
  resp = validate(api_key)
45
  if ("error" in resp.json()):
46
- st.info("Your API Token is incorrect! Try again.")
47
  else:
 
48
  os.environ["OPENAI_API_KEY"] = api_key
49
- index = create_index()
50
 
51
  st.write("---")
52
  input_text = st.text_area("Ask your question")
53
 
54
  if input_text is not None:
55
  if st.button("Ask"):
56
- st.info("Your query: \n" + input_text)
57
  with st.spinner("Processing your query..."):
58
- response = get_response(index, input_text)
59
- print(response)
60
-
61
- st.success(response)
62
 
63
- st.write("---")
 
 
 
2
 
3
  import requests
4
  import streamlit as st
5
+ from streamlit_chat import message
6
 
7
  from models import bloom
8
+ from utils.util import *
9
 
10
  st.title("Welcome to RegBotBeta")
11
  st.header("Powered by `LlamaIndex🦙` and `OpenAI API`")
12
 
13
+ if 'messages' not in st.session_state:
14
+ st.session_state.messages = []
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
15
 
16
  api_key = st.text_input("Enter your OpenAI API key here:", type="password")
17
  if api_key:
18
  resp = validate(api_key)
19
  if ("error" in resp.json()):
20
+ st.info("Invalid Token! Try again.")
21
  else:
22
+ st.info("Success")
23
  os.environ["OPENAI_API_KEY"] = api_key
24
+ index = create_index(bloom)
25
 
26
  st.write("---")
27
  input_text = st.text_area("Ask your question")
28
 
29
  if input_text is not None:
30
  if st.button("Ask"):
31
+ st.session_state.messages.append(('User', input_text))
32
  with st.spinner("Processing your query..."):
33
+ bot_response = get_response(index, input_text)
34
+
35
+ st.session_state.messages.append(('Bot', bot_response))
 
36
 
37
+ # Display previous messages
38
+ for sender, msg in st.session_state.messages[::-1]:
39
+ message(msg, is_user = (sender=='User'))
models/bloom.py CHANGED
@@ -14,9 +14,9 @@ from utils.customLLM import CustomLLM
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
@@ -28,20 +28,20 @@ chunk_overlap_ratio = 0.2
28
  prompt_helper = PromptHelper(context_window, num_output, chunk_overlap_ratio)
29
 
30
  # create a pipeline
31
- pl = pipeline(
32
- model=model,
33
- tokenizer=tokenizer,
34
- task="text-generation",
35
- # device=0, # GPU device number
36
- # max_length=512,
37
- do_sample=True,
38
- top_p=0.95,
39
- top_k=50,
40
- temperature=0.7
41
- )
42
 
43
  # define llm
44
- llm_predictor = LLMPredictor(llm=CustomLLM(model_pipeline=pl))
45
  service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
46
 
47
  def prepare_data(file_path:str):
 
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
 
28
  prompt_helper = PromptHelper(context_window, num_output, chunk_overlap_ratio)
29
 
30
  # create a pipeline
31
+ # pl = pipeline(
32
+ # model=model,
33
+ # tokenizer=tokenizer,
34
+ # task="text-generation",
35
+ # # device=0, # GPU device number
36
+ # # max_length=512,
37
+ # do_sample=True,
38
+ # top_p=0.95,
39
+ # top_k=50,
40
+ # temperature=0.7
41
+ # )
42
 
43
  # define llm
44
+ llm_predictor = LLMPredictor(llm=CustomLLM())
45
  service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
46
 
47
  def prepare_data(file_path:str):
utils/customLLM.py CHANGED
@@ -1,14 +1,26 @@
1
  from typing import Any, List, Mapping, Optional
2
 
3
  from langchain.llms.base import LLM
4
- from transformers import Pipeline
5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6
 
7
  class CustomLLM(LLM):
8
- pipeline = None
9
- # Create the pipeline for question answering
10
- def __init__(self, model_pipeline: Pipeline):
11
- self.pipeline = model_pipeline
12
 
13
  def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
14
  prompt_length = len(prompt)
 
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
+ model_name = "bigscience/bloom-560m"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
8
+ model = AutoModelForCausalLM.from_pretrained(model_name, config='T5Config')
9
+
10
+ pl = pipeline(
11
+ model=model,
12
+ tokenizer=tokenizer,
13
+ task="text-generation",
14
+ # device=0, # GPU device number
15
+ # max_length=512,
16
+ do_sample=True,
17
+ top_p=0.95,
18
+ top_k=50,
19
+ temperature=0.7
20
+ )
21
 
22
  class CustomLLM(LLM):
23
+ pipeline = pl
 
 
 
24
 
25
  def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
26
  prompt_length = len(prompt)