tien314/hscode_classification

1import streamlit as st2import bm25s3from operator import itemgetter4import os5import re6import pandas as pd7from langchain_groq import ChatGroq8from langchain_core.prompts import ChatPromptTemplate, HumanMessagePromptTemplate9from langchain.docstore.document import Document10 11@st.cache_data12def load_data():13    df = pd.read_csv("cleaned_list.csv",header = None)14    df.columns = ['document']15    corpus = [doc for doc in df['document'].to_list()]16 17    retriever = bm25s.BM25(corpus=corpus)18    retriever.index(bm25s.tokenize(corpus))19 20    return retriever21 22# def extract_hscode(text):23#     match = re.search(r'hs_code:\s*(\d+)', text)24#     if match:25#         return match.group(1)26#     return None27 28# df2 = pd.read_csv("hscode_main.csv")29# new_col = [len(str(code))for code in df2['hs_code'].to_list()]30# df2['len'] = new_col31 32# new_hscode = [str(code) for code in df2['hs_code']]33 34# for i in range(len(new_col)):35#     if new_col[i]==5:36#         new_hscode[i] = '0'+ new_hscode[i]37# df2['hs_code'] = new_hscode38# df2=df2.drop(columns='len')39 40# if 'retriever' not in st.session_state:41#     st.session_state.retriever = None42 43# if st.session_state.retriever is None:44#     st.session_state.retriever = load_data()45 46 47# sentence = st.text_input("please enter description:")48 49# if sentence !='':50#     results,_ = st.session_state.retriever.retrieve(bm25s.tokenize(sentence), k=5)51#     doc = [d for d in results]52#     hscodes = [extract_hscode(item) for item in doc[0]]53#     for code in hscodes:54#         if len(code)==5:55#             code = '0'+ code56 57#         filter_df = df2[df2['hs_code']==code]58#         answer = filter_df['description'].iloc[0]59#         st.write("Hscode:",code)60#         st.write("Description:",answer.lower())61 62def load_model():63    prompt = ChatPromptTemplate.from_messages([64        HumanMessagePromptTemplate.from_template(65        f"""66        Extract the appropriate 6-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.67        Only return the HS Code as a 6-digit number .68        Example: 12345669        Context: {{context}}70        Description: {{description}}71        Answer:72        """73        )74    ])75    76 77    #device = "cuda" if torch.cuda.is_available() else "cpu"78    79    #llm = OllamaLLM(model="gemma2", temperature=0, device=device)80    #api_key = "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"81    api_key = "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"82    llm = ChatGroq(model = "llama-3.1-70b-versatile", temperature = 0,api_key = api_key)83    chain = prompt|llm84    return chain85 86def process_input(sentence):87    docs, _ = st.session_state.retriever.retrieve(bm25s.tokenize(sentence), k=15)88    documents =[]89    for doc in docs[0]:90        documents.append(Document(doc)) 91    return documents92    93if 'retriever' not in st.session_state:94    st.session_state.retriever = None95 96if 'chain' not in st.session_state:97    st.session_state.chain = None98    99if st.session_state.retriever is None:100    st.session_state.retriever = load_data()101 102if st.session_state.chain is None:103    st.session_state.chain = load_model()104    105sentence = st.text_input("please enter description:")106 107if sentence !='':108    documents = process_input(sentence)109    hscode = st.session_state.chain.invoke({'context': documents,'description':sentence})110    st.write("answer:",hscode.content)