File size: 1,541 Bytes
5020dd3 |
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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import streamlit as st
import joblib
from sentence_transformers import SentenceTransformer
# Load the pre-trained embedding model
@st.cache_resource # Cache the embedding model to save loading time
def load_embedding_model():
return SentenceTransformer('neuml/pubmedbert-base-embeddings')
# Load the MLP model
@st.cache_resource # Cache the loaded model
def load_mlp_model():
with open("MLP.pkl", "rb") as file:
return joblib.load(file)
# Embed text
def get_embeddings(title, abstract, embedding_model):
# Concatenate title and abstract
combined_text = title + " " + abstract
return embedding_model.encode(combined_text)
# Main Streamlit app
def main():
st.title("MLP Predictor for Titles and Abstracts")
# Input fields
title = st.text_input("Enter the Title:")
abstract = st.text_area("Enter the Abstract:")
# Load models
embedding_model = load_embedding_model()
mlp_model = load_mlp_model()
# Predict button
if st.button("Predict Label"):
if title.strip() == "" or abstract.strip() == "":
st.error("Both Title and Abstract are required!")
else:
# Get embeddings
embeddings = get_embeddings(title, abstract, embedding_model)
# Make prediction
prediction = mlp_model.predict([embeddings])[0] # Input should be a 2D array
# Display result
st.success(f"The predicted label is: {prediction}")
if __name__ == "__main__":
main()
|