fixed some bugs

app.py

@@ -2,62 +2,38 @@ import os
 
 import requests
 import streamlit as st
+from streamlit_chat import message
 
 from models import bloom
+from utils.util import *
 
 st.title("Welcome to RegBotBeta")
 st.header("Powered by `LlamaIndex🦙` and `OpenAI API`")
 
-def validate(token: str):
-    api_endpoint = "https://api.openai.com/v1/chat/completions"
-    api_key = token
-
-    headers = {
-        "Content-Type" : "application/json",
-        "Authorization": f"Bearer {api_key}"
-    }
-
-    messages = [
-        {"role": "user", "content": "Say this is a test!"}
-    ]
-
-    data = {
-        "model": "gpt-3.5-turbo",
-        "messages": messages
-    }
-
-    response = requests.post(api_endpoint, json=data, headers=headers)
-    return response
-
-def create_index():
-    index = bloom.initialize_index("bloomLlama")
-    return index
-
-def get_response(vector_index, query_str):
-    query_engine = vector_index.as_query_engine()
-    response = query_engine.query(query_str)
-    return response
-
+if 'messages' not in st.session_state:
+    st.session_state.messages = []
 
 api_key = st.text_input("Enter your OpenAI API key here:", type="password")
 if api_key:
     resp = validate(api_key)
     if ("error" in resp.json()):
-        st.info("Your API Token is incorrect! Try again.")
+        st.info("Invalid Token! Try again.")
     else:
+        st.info("Success")
         os.environ["OPENAI_API_KEY"] = api_key
-        index = create_index()
+        index = create_index(bloom)
 
 st.write("---")
 input_text = st.text_area("Ask your question")
 
 if input_text is not None:
     if st.button("Ask"):
-        st.info("Your query: \n" + input_text)
+        st.session_state.messages.append(('User', input_text))
         with st.spinner("Processing your query..."):
-            response = get_response(index, input_text)
-            print(response)
-
-        st.success(response)
+            bot_response = get_response(index, input_text)
+        
+        st.session_state.messages.append(('Bot', bot_response))
 
-        st.write("---")
+# Display previous messages
+for sender, msg in st.session_state.messages[::-1]:
+    message(msg, is_user = (sender=='User'))

models/bloom.py

@@ -14,9 +14,9 @@ from utils.customLLM import CustomLLM
 load_dotenv()
 
 # get model
-model_name = "bigscience/bloom-560m" 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, config='T5Config')
+# model_name = "bigscience/bloom-560m" 
+# tokenizer = AutoTokenizer.from_pretrained(model_name)
+# model = AutoModelForCausalLM.from_pretrained(model_name, config='T5Config')
 
 # define prompt helper
 # set maximum input size
@@ -28,20 +28,20 @@ chunk_overlap_ratio = 0.2
 prompt_helper = PromptHelper(context_window, num_output, chunk_overlap_ratio)
 
 # create a pipeline
-pl = pipeline(
-    model=model,
-    tokenizer=tokenizer,
-    task="text-generation",
-    # device=0, # GPU device number
-    # max_length=512,
-    do_sample=True,
-    top_p=0.95,
-    top_k=50,
-    temperature=0.7
-)
+# pl = pipeline(
+#     model=model,
+#     tokenizer=tokenizer,
+#     task="text-generation",
+#     # device=0, # GPU device number
+#     # max_length=512,
+#     do_sample=True,
+#     top_p=0.95,
+#     top_k=50,
+#     temperature=0.7
+# )
 
 # define llm
-llm_predictor = LLMPredictor(llm=CustomLLM(model_pipeline=pl))
+llm_predictor = LLMPredictor(llm=CustomLLM())
 service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
 
 def prepare_data(file_path:str):

utils/customLLM.py

@@ -1,14 +1,26 @@
 from typing import Any, List, Mapping, Optional
 
 from langchain.llms.base import LLM
-from transformers import Pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
+model_name = "bigscience/bloom-560m" 
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, config='T5Config')
+
+pl = pipeline(
+    model=model,
+    tokenizer=tokenizer,
+    task="text-generation",
+    # device=0, # GPU device number
+    # max_length=512,
+    do_sample=True,
+    top_p=0.95,
+    top_k=50,
+    temperature=0.7
+)
 
 class CustomLLM(LLM):
-    pipeline = None
-    # Create the pipeline for question answering
-    def __init__(self, model_pipeline: Pipeline):
-        self.pipeline = model_pipeline
+    pipeline = pl
 
     def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
         prompt_length = len(prompt)