Update app.py

app.py

@@ -1,16 +1,39 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import re
+from transformers import AutoTokenizer, T5EncoderModel
+import torch.nn as nn
 
-# Modell laden
-model = AutoModelForSequenceClassification.from_pretrained("pepegiallo/flan-t5-base_ner")
+# Klassendefinition aus dem Training
+class FlanT5Classifier(nn.Module):
+    def __init__(self, base_model_name="google/flan-t5-base", num_labels=4):
+        super().__init__()
+        self.encoder = T5EncoderModel.from_pretrained(base_model_name)
+        self.dropout = nn.Dropout(0.1)
+        self.classifier = nn.Linear(self.encoder.config.d_model, num_labels)
+
+    def forward(self, input_ids, attention_mask=None):
+        encoder_outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
+        pooled = encoder_outputs.last_hidden_state[:, 0]
+        logits = self.classifier(self.dropout(pooled))
+        return {"logits": logits}
+
+# Tokenizer laden
 tokenizer = AutoTokenizer.from_pretrained("pepegiallo/flan-t5-base_ner")
+
+# Modell instanziieren und Token-Embeddings anpassen
+model = FlanT5Classifier()
+model.encoder.resize_token_embeddings(len(tokenizer))
+
+# Gewichte laden
+state_dict = torch.load("pytorch_model.bin", map_location="cpu")
+model.load_state_dict(state_dict)
 model.eval()
 
+# ID-Zuordnung
 id2label = {0: "LOC", 1: "ORG", 2: "PER", 3: "O"}
 
-# Hilfsfunktionen
+# Tokenizer-Funktionen
 def custom_tokenize(text):
     return re.findall(r"\w+|[^\w\s]", text, re.UNICODE)
 
@@ -22,6 +45,7 @@ def custom_detokenize(tokens):
         text += token
     return text
 
+# Klassifikationsfunktion
 def classify_tokens(text):
     tokens = custom_tokenize(text)
     results = []
@@ -32,14 +56,14 @@ def classify_tokens(text):
 
         inputs = tokenizer(prompt, return_tensors="pt", padding="max_length", truncation=True, max_length=128)
         with torch.no_grad():
-            logits = model(**inputs).logits
+            logits = model(**inputs)["logits"]
             pred_id = torch.argmax(logits, dim=-1).item()
             label = id2label[pred_id]
 
         results.append((tokens[i], label))
     return results
 
-# Gradio-UI definieren
+# Gradio UI
 demo = gr.Interface(
     fn=classify_tokens,
     inputs=gr.Textbox(lines=3, placeholder="Enter a sentence..."),