Update app.py

app.py

@@ -3,15 +3,9 @@ import requests
 import os
 from streamlit_chat import message
 import random
+import sentence-transformers
+import nltk
 
-@st.cache
-def query(payload):
-    api_token = os.getenv("api_token")
-    model_id = "deepset/roberta-base-squad2"
-    headers = {"Authorization": f"Bearer {api_token}"}
-    API_URL = f"https://api-inference.huggingface.co/models/{model_id}"
-    response = requests.post(API_URL, headers=headers, json=payload)
-    return response.json(), response
 
 
 context = "To extract information from documents, use sentence similarity task. To classify sentiments, use text classification task. To do sentiment analysis, use text classification task. To detect masks from images, use object detection task. To extract name or address from documents use token classification task. To extract name or address from invoices, use token classification task. To build voice enabled applications, you can use automatic speech recognition task. You can retrieve information from documents using sentence similarity task. You can summarize papers using summarization task. You can convert text to speech using text-to-speech task. To detect language spoken in an audio, you can use audio classification task. To detect emotion in an audio, you can use audio classification task. To detect commands in an audio, you can use audio classification task. To decompose sounds in a recording, use audio-to-audio task. To answer questions from a document, you can use question answering task. To answer FAQs from your customers, you can use question answering task. To see if a text is grammatically correct, you can use text classification task. To augment your training data, you can use text classification task. To detect pedestrians, you can use object detection task."
@@ -34,6 +28,51 @@ link_dict = {
     "object detection": "https://huggingface.co/tasks/object-detection"}
     
 
+
+
+corpus = []
+sentence_count = []
+
+for sent in context.split("."):
+
+    sentences = nltk.tokenize.sent_tokenize(str(sent), language='english')
+    sentence_count.append(len(sentences))
+    for _,s in enumerate(sentences):
+        corpus.append(s)
+        
+        
+        
+corpus_embeddings = np.load('task_embeddings_msmarco-distilbert-base-v4.npy')
+corpus_embeddings.shape
+
+
+
+
+def find_sentences(query):
+    query_embedding = model.encode(query)
+    hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=5)
+    hits = hits[0]
+
+    for hit in hits:
+        corpus_id = hit['corpus_id']
+        print(corpus[corpus_id], "(Score: {:.4f})".format(hit['score']))
+        # Find source document based on sentence index
+        count = 0
+        for idx, c in enumerate(sentence_count):
+            count+=c
+            if (corpus_id > count-1):
+                continue
+            else:
+                doc = corpus[idx]
+                print(f"Document: {doc}, {count}")
+                break
+
+
+
+
+
+
+
 message_history = [{"text":"Let's find out the best task for your use case! Tell me about your use case :)", "is_user":False}]
 
 st.subheader("If you don't know how to build your machine learning product for your use case, Taskmaster is here to help you! 🪄✨")
@@ -46,21 +85,15 @@ placeholder = st.empty()  # placeholder for latest message
 
 input = st.text_input("Ask me 🤗")
 if input:
+
     message_history.append({"text":input, "is_user" : True})
 
-    data, resp = query(
-    {
-        "inputs": {
-            "question": input,
-            "context": context,
-        }
-    }
-    )
+    model_answer = find_sentences("How can I extract information from invoices?")
 
 
     if resp.status_code == 200:
     
-        model_answer = data["answer"]
+
         key_exists = False
         for key in link_dict:
             if key in model_answer: