divided py files and added more error catching

app.py

@@ -2,68 +2,7 @@ import os
 import random
 import time
 import gradio as gr
-import subprocess
-import requests
-from langchain.chains.question_answering import load_qa_chain
-from langchain.text_splitter import CharacterTextSplitter
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.docstore.document import Document
-from langchain.document_loaders import TextLoader
-from langchain.vectorstores import FAISS
-from langchain import HuggingFaceHub
-
-
-class DocumentChatbot:
-    def __init__(self):
-        self.llm = None
-        self.chain = None
-        self.embeddings = None
-        self.metadata = {"source": "internet"}
-        self.init_mes = ["According to the document, ", "Based on the text, ", "I think, ", "According to the text, ", "Based on the document you provided, "]
-        
-    def load_model(self, api_key):
-        os.environ["HUGGINGFACEHUB_API_TOKEN"] = api_key
-        result = subprocess.run(["curl", "https://huggingface.co/api/whoami-v2", "-H", f"Authorization: Bearer {api_key}"], capture_output=True).stdout.decode()
-        self.llm = HuggingFaceHub(repo_id="google/flan-t5-large", model_kwargs={"temperature":0, "max_length":512})
-        self.chain = load_qa_chain(self.llm, chain_type="stuff")
-        self.embeddings = HuggingFaceEmbeddings()
-        if result == '{"error":"Invalid username or password."}':
-          return  "Invalid API token"
-        else:
-          return "HF Token successfully registered"
-
-
-    def respond(self, text_input, question, chat_history):
-        if text_input.startswith("http"):
-            response = requests.get(text_input)
-            text_var = response.text
-        else:
-            text_var = text_input
-
-        time.sleep(0.5)
-
-        documents = [Document(page_content=text_var, metadata=self.metadata)]
-        text_splitter = CharacterTextSplitter(chunk_size=750, chunk_overlap=0)
-        docs = text_splitter.split_documents(documents)
-
-        if self.llm is None:
-            raise ValueError("Model not loaded")
-
-        db = FAISS.from_documents(docs, self.embeddings)
-        query = question
-
-        start_time = time.monotonic()
-        try:
-            docs = db.similarity_search(query)
-            answer = self.chain.run(input_documents=docs, question=query, max_execution_time=5)
-        except ValueError as e:
-            answer = f"An error occurred: {str(e)}"
-
-        bot_message = random.choice(self.init_mes) + answer + "."
-        chat_history.append((question, bot_message))
-        time.sleep(1)
-        return "", chat_history
-
+from document_chatbot import DocumentChatbot
 
 document_chatbot = DocumentChatbot()
 
@@ -71,18 +10,19 @@ with gr.Blocks() as demo:
     title = """<p><h1 align="center" style="font-size: 36px;">Talk with your document</h1></p>"""
     gr.HTML(title)
     with gr.Row():
-      text_input = gr.Textbox(label="Enter text or URL to text file")
-      with gr.Column():
-        api_key_input = gr.Textbox(label="Enter HF Token to load the model")
-        api_key_input.submit(document_chatbot.load_model, inputs=api_key_input, outputs=api_key_input)
-        chatbot = gr.Chatbot()
+        text_input = gr.Textbox(label="Enter text or URL to text file")
+        with gr.Column():
+            with gr.Row():
+                api_key_input = gr.Textbox(label="Enter HF Token to load the model")
+                api_key_input.submit(document_chatbot.load_token, inputs=api_key_input, outputs=api_key_input)
+                picked_model = gr.Dropdown(["google/flan-t5-large", "google/flan-t5-base","google/flan-t5-small"], label="Models", info="I'd recommend choosing  the first one")
+                picked_model.change(document_chatbot.load_model, picked_model)
+            chatbot = gr.Chatbot()
+
+            q_input = gr.Textbox(label="Please write your question")
+            clear = gr.Button("Clear")
+            q_input.submit(document_chatbot.respond, [text_input, q_input, chatbot], [q_input, chatbot])
+            clear.click(lambda: None, None, chatbot, queue=False)
         
-        q_input = gr.Textbox(label="Please write your question")
-        clear = gr.Button("Clear")
-        q_input.submit(document_chatbot.respond, [text_input, q_input, chatbot], [q_input, chatbot])
-        clear.click(lambda: None, None, chatbot, queue=False)
-      
-
-
         
-demo.launch(debug=True)
+demo.launch(debug=True)

document_chatbot.py

@@ -0,0 +1,75 @@
+import os
+import random
+import time
+import subprocess
+import requests
+from langchain.chains.question_answering import load_qa_chain
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.docstore.document import Document
+from langchain.document_loaders import TextLoader
+from langchain.vectorstores import FAISS
+from langchain import HuggingFaceHub
+
+
+class DocumentChatbot:
+    
+    def __init__(self):
+        self.llm = None
+        self.chain = None
+        self.embeddings = None
+        self.metadata = {"source": "internet"}
+        self.init_mes = ["According to the document, ", "Based on the text, ", "I think, ", "According to the text, ", "Based on the document you provided, "]
+        
+    def load_token(self, api_key):
+        if api_key[:2] == "hf":
+            os.environ["HUGGINGFACEHUB_API_TOKEN"] = api_key
+            result = subprocess.run(["curl", "https://huggingface.co/api/whoami-v2", "-H", f"Authorization: Bearer {api_key}"], capture_output=True).stdout.decode()
+        if result == '{"error":"Invalid username or password."}':
+            return  "Invalid API token"
+        else:
+            return "HF Token successfully registered"
+        
+    def load_model(self, model_name):
+        self.llm = HuggingFaceHub(repo_id=model_name, model_kwargs={"temperature":0, "max_length":512})
+        self.chain = load_qa_chain(self.llm, chain_type="stuff")
+        self.embeddings = HuggingFaceEmbeddings()
+        return f"Successfully loaded {model_name}"
+
+
+    def respond(self, text_input, question, chat_history):
+        if text_input.startswith("http"):
+            response = requests.get(text_input)
+            text_var = response.text
+            if text_var is None:
+                raise ValueError("No document is given")
+        else:
+            text_var = text_input
+            
+
+
+        time.sleep(0.5)
+
+        documents = [Document(page_content=text_var, metadata=self.metadata)]
+        text_splitter = CharacterTextSplitter(chunk_size=750, chunk_overlap=0)
+        docs = text_splitter.split_documents(documents)
+
+        if self.llm is None:
+            raise ValueError("Model not loaded")
+            
+
+
+        db = FAISS.from_documents(docs, self.embeddings)
+        query = question
+        
+        try:
+            docs = db.similarity_search(query)
+            answer = self.chain.run(input_documents=docs, question=query)
+            bot_message = random.choice(self.init_mes) + answer + "."
+        except ValueError as e:
+            bot_message = f"An error occurred: {str(e)}"
+        chat_history.append((question, bot_message))
+
+        time.sleep(1)
+        
+        return "", chat_history