File size: 710 Bytes
9abea4e
 
1434337
 
4c2a969
 
 
 
9abea4e
274b354
 
408dd7e
 
 
 
 
5b26a96
 
 
 
 
 
 
 
 
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
import streamlit as st

INDEX_DIR = "data/index"
STATEMENTS_PATH = "data/statements.txt"

RETRIEVER_MODEL = "sentence-transformers/msmarco-distilbert-base-tas-b"
RETRIEVER_MODEL_FORMAT = "sentence_transformers"
RETRIEVER_TOP_K = 5

# In HF Space, we use microsoft/deberta-v2-xlarge-mnli
# for local testing, a smaller model is better
try:
    NLI_MODEL = st.secrets["NLI_MODEL"]
except:
    NLI_MODEL = "valhalla/distilbart-mnli-12-1"
print(f"Used NLI model: {NLI_MODEL}")


# In HF Space, we use google/flan-t5-large
# for local testing, a smaller model is better
try:
    PROMPT_MODEL = st.secrets["PROMPT_MODEL"]
except:
    PROMPT_MODEL = "google/flan-t5-small"
print(f"Used Prompt model: {PROMPT_MODEL}")