add trained model

app.py

@@ -4,7 +4,7 @@ from inference import classify, load_model  # Replace with your actual module
 
 @st.cache_resource
 def get_model():
-    return load_model()
+    return load_model("kalinin-a-i/ml2-hw4", "class2name.joblib")
 
 def get_arxiv_article_info(url):
     """Extracts title and abstract from an arXiv article link."""
@@ -16,7 +16,7 @@ def get_arxiv_article_info(url):
     return None, None
 
 # Load model once
-model = get_model()
+pipe, class2name = get_model()
 
 st.title("ArXiv Article Classifier")
 
@@ -41,6 +41,6 @@ elif input_method == "Manual Input":
 
 # Classification and output
 if title and abstract:
-    category = classify(title, abstract)
+    category = classify(pipe, class2name, title, abstract)
     st.write(f"### Title: {title}")
     st.write(f"**Predicted Category:** {category}")

inference.py

@@ -1,6 +1,24 @@
-def load_model():
-    pass
+from transformers import  AutoTokenizer
+from transformers import AutoModelForSequenceClassification
+from transformers import pipeline, Pipeline
+from joblib import load
 
-def classify(title, abstract):
-    return "physics"
+
+def load_model(path2chkpt: str, path2mapping: str):
+    model = AutoModelForSequenceClassification.from_pretrained("/home/jupyter/datasphere/project/hw4_nlp_ops/weights_20_classes/checkpoint-4500")
+    tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-cased")
+
+    pipe = pipeline("text-classification", 
+                   model=model, 
+                   tokenizer=tokenizer)
+
+    class2name = load(path2mapping)
+    return pipe, class2name
+
+
+def classify(pipe: Pipeline, class2name: dict[str, str], title: str, abstract: str):
+    inputs = ".".join([title, abstract])
+    class_code = pipe(inputs)[0]["label"]
+    
+    return class2name[class_code]
 

requirements.txt

@@ -1,2 +1,3 @@
 streamlit
-arxiv
+arxiv
+joblib