File size: 1,343 Bytes
7b0074a
 
 
 
 
f1aebfb
7b0074a
de3cc67
a735075
590a02d
de3cc67
 
 
 
 
 
 
590a02d
57de215
638c04e
57de215
7b0074a
 
57de215
7b0074a
 
 
 
57de215
 
 
b908f48
0cc9ff0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_community.document_loaders import DirectoryLoader
from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from langchain_community.vectorstores import Chroma
import streamlit as st

lang = st.radio("Choose Language:",
               ["English", "Chinese"],
               captions= ['Support in English', 'Chinese QnA'])
if lang == "English":
    embedding_function = SentenceTransformerEmbeddings(model_name="all-mpnet-base-v2")
    chdb = Chroma(persist_directory="./chroma_db_info", embedding_function=embedding_function)
else:
    embedding_function = SentenceTransformerEmbeddings(model_name="shibing624/text2vec-base-chinese")
    chdb = Chroma(persist_directory="./chroma_db_ch", embedding_function=embedding_function)
   
query = st.text_input("Enter a query about IC design:")
disnum = 3
if query:
    docs = chdb.similarity_search_with_score(query)
    docnum = len(docs)
    index = 0
    ret = f"Query:{query}\n"
    for ii in range(docnum):    
        doc = docs[ii][0]
        score = docs[ii][1]
        ret += f"Return {index} ({score:.4f}) :\n{doc.page_content}\n"
        index += 1
        if index > disnum:
            break
    st.text(ret)