Update app.py

app.py

@@ -14,6 +14,9 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # Load models
 editorial_model = "PleIAs/Estienne"
 bibliography_model = "PleIAs/Bibliography-Formatter"
+bibliography_style = "PleIAs/Bibliography-Classifier"
+
+tokenizer = AutoTokenizer.from_pretrained(editorial_model, model_max_length=512)
 
 editorial_classifier = pipeline(
     "token-classification", model=editorial_model, aggregation_strategy="simple", device=device
@@ -22,7 +25,7 @@ bibliography_classifier = pipeline(
     "token-classification", model=bibliography_model, aggregation_strategy="simple", device=device
 )
 
-tokenizer = AutoTokenizer.from_pretrained(editorial_model, model_max_length=512)
+style_classifier = pipeline("text-classification", model=bibliography_style, tokenizer=tokenizer, device=device)
 
 # Helper functions
 def preprocess_text(text):
@@ -172,6 +175,10 @@ class CombinedProcessor:
         bibliography_entries = editorial_df[editorial_df['entity_group'] == 'bibliography']['word'].tolist()
         
         bibtex_entries = []
+
+        for entry in bibliography_entries:
+            style = style_classifier(entry, truncation=True, padding=True, top_k=None)
+            print(style)
         
         for entry in bibliography_entries:
             print(entry)