Update app.py

app.py

@@ -7,8 +7,8 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 st.title(':blue[_SnapCode_]')
 
 with st.spinner('Model is loading'):
-    tokenizer = AutoTokenizer.from_pretrained("vishnun/codenlbert")
-    model = AutoModelForSequenceClassification.from_pretrained("vishnun/codenlbert")
+    tokenizer = AutoTokenizer.from_pretrained("vishnun/codenlbert-sm")
+    model = AutoModelForSequenceClassification.from_pretrained("vishnun/codenlbert-sm")
 
 st.success('Model loaded')
 
@@ -29,7 +29,7 @@ if uploaded_file is not None:
     ocr_class = [classify_text(x) for x in ocr_list]
     idx = []
     for i in range(len(ocr_class)):
-        if ocr_class[i] == 'CODE':
+        if ocr_class[i].upper() == 'CODE':
             idx.append(ocr_list[i])