Update app.py

app.py

@@ -1,108 +1,68 @@
 import streamlit as st
-from langchain_community.retrievers import BM25Retriever
 import pandas as pd
+import bm25s
+from bm25s.hf import BM25HF
+from langchain_core.prompts import ChatPromptTemplate, HumanMessagePromptTemplate
 from langchain.docstore.document import Document
-from langchain.text_splitter import CharacterTextSplitter
-from operator import itemgetter
-from langchain_core.prompts import PromptTemplate
-from langchain_groq import ChatGroq
-from langchain.chains.question_answering import load_qa_chain
+import torch
 import os
+from huggingface_hub import login
+from langchain_groq import ChatGroq
 
 
-@st.cache_data
+@st.cache_resource
 def load_data():
-    df = pd.read_csv("trained.csv")
-    df = df.drop(columns = ['Unnamed: 0','hs_code_2','hs_code_4'])
-    documents = []
-
-    for index, row in df.iterrows():
-        text = row['full_description']
-        hs_code = row['hs_code']
-        documents.append(Document(page_content=text, metadata={'hs_code': hs_code}))
-
-    splitter = CharacterTextSplitter(
-        chunk_size=100,  
-        chunk_overlap=0,
-        separator = ' '
-    )
-
-    split_documents = []
-    for doc in documents:
-        chunks = splitter.split_text(doc.page_content)
-        #remove chunk split word
-        word_chunks = []
-        current_chunk = []
-        
-        for chunk in chunks:
-            words = chunk.split()
-            for word in words:
-                if len(' '.join(current_chunk + [word])) <=100:
-                    current_chunk.append(word)
-                else:
-                    word_chunks.append(' '.join(current_chunk))
-                    current_chunk = [word]
-        if current_chunk:
-            word_chunks.append(' '.join(current_chunk))
-            
-        split_documents.append(Document(page_content=word_chunks[0], metadata=doc.metadata))
-
-
-    docs = []
-    for doc in split_documents:
-        metadata = doc.metadata
-        metadata_str = str(metadata).strip('{}')
-        page = doc.page_content
-        docs.append([metadata_str + " " + page])
-
-
-    cleaned_list = [item.replace('"','').replace("'",'') for items in docs for item in items]
-    retriever = BM25Retriever.from_texts(cleaned_list)
-    retriever.k = 5
+    
+    retriever = BM25HF.load_from_hub(
+    "tien314/hscode8", load_corpus=True, mmap=True)
     return retriever
 
-
-def load_llm():
-
-    api_key2 = "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"
-
-    llm2 = ChatGroq(model = "llama-3.1-70b-versatile", temperature = 0,api_key = api_key2)
-    return llm2
-
-
-def predict(sentence,retriever,llm2):
-    sentence = sentence.lower()
-    context = retriever.get_relevant_documents(sentence)
-    #print("context:",context)
-    template2 = """
-    You are an expert in HS Code classification. 
-    Based on the provided product description, accurately determine and return only one 6-digit HS Code that best matches the description.
-    Always return the HS Code as a 6-digit number only.
-    example: 123456
-    Context:\n {context} \n
-    Description:\n {description} \n
-    Answer:
-    """
-    prompt2 = PromptTemplate(template=template2, input_variables=['context','description'])
-    chain = load_qa_chain(llm2, chain_type = 'stuff', prompt = prompt2)
-    response = chain.invoke({'input_documents': context, 'description':sentence})
-    answer = response.get("output_text")
-    return answer
-
-
+def load_model():
+    prompt = ChatPromptTemplate.from_messages([
+        HumanMessagePromptTemplate.from_template(
+        f"""
+        Extract the appropriate 8-digit HS Code base on the product description and retrieved document by thoroughly analyzing its details and utilizing a reliable and up-to-date HS Code database for accurate results.
+        Only return the HS Code as a 8-digit number .
+        Example: 1234567878
+        Context: {{context}}
+        Description: {{description}}
+        Answer:
+        """
+        )
+    ])
+    
+
+    #device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    #llm = OllamaLLM(model="gemma2", temperature=0, device=device)
+    #api_key = "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"
+    api_key = "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"
+    llm = ChatGroq(model = "llama-3.1-70b-versatile", temperature = 0,api_key = api_key)
+    chain = prompt|llm
+    return chain
+
+def process_input(sentence):
+    docs, _ = st.session_state.retriever.retrieve(bm25s.tokenize(sentence), k=15)
+    documents =[]
+    for doc in docs[0]:
+        documents.append(Document(doc['text'])) 
+    return documents
+    
 if 'retriever' not in st.session_state:
     st.session_state.retriever = None
-if 'llm' not in st.session_state:
-    st.session_state.llm = None
 
+if 'chain' not in st.session_state:
+    st.session_state.chain = None
+    
 if st.session_state.retriever is None:
     st.session_state.retriever = load_data()
 
-if st.session_state.llm is None:
-    st.session_state.llm = load_llm()
-
+if st.session_state.chain is None:
+    st.session_state.chain = load_model()
+    
 sentence = st.text_input("please enter description:")
 
 if sentence !='':
-    answer = predict(sentence,st.session_state.retriever,st.session_state.llm )
-    st.write("answer:",answer)
+    documents = process_input(sentence)
+    hscode = st.session_state.chain.invoke({'context': documents,'description':sentence})
+    st.write("answer:",hscode.content)