pleonova commited on
Commit
b52e509
·
verified ·
1 Parent(s): 539f2fe

Change to FastAPI

Browse files
Files changed (1) hide show
  1. app.py +8 -34
app.py CHANGED
@@ -1,37 +1,11 @@
1
- import streamlit as st
2
  from transformers import pipeline
3
 
4
- # Initialize the zero-shot classification pipeline
5
- @st.cache_resource
6
- def load_classifier():
7
- return pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
8
 
9
- classifier = load_classifier()
10
-
11
- # Streamlit app interface
12
- st.title("Webpage Subject Classifier")
13
- st.write(
14
- """
15
- Enter the text content of a webpage below to classify it into one of the following subjects:
16
- - Mathematics
17
- - Language Arts
18
- - Social Studies
19
- - Science
20
- """
21
- )
22
-
23
- # Text input area
24
- text_input = st.text_area("Paste the webpage content here:", height=200)
25
-
26
- # Define the candidate labels
27
- labels = ["Mathematics", "Language Arts", "Social Studies", "Science"]
28
-
29
- # Perform classification when the "Classify" button is clicked
30
- if st.button("Classify"):
31
- if text_input.strip():
32
- with st.spinner("Classifying..."):
33
- result = classifier(text_input, labels)
34
- predicted_label = result["labels"][0]
35
- st.success(f"The predicted subject is: **{predicted_label}**")
36
- else:
37
- st.warning("Please enter some text to classify.")
 
1
+ from fastapi import FastAPI
2
  from transformers import pipeline
3
 
4
+ app = FastAPI()
5
+ classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
 
 
6
 
7
+ @app.post("/predict")
8
+ async def predict(data: dict):
9
+ labels = ["Mathematics", "Language Arts", "Social Studies", "Science"]
10
+ result = classifier(data["text"], labels)
11
+ return {"label": result["labels"][0]}