app

app.py

@@ -0,0 +1,13 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+from pipeline import DiscoursePipeline  # ton code custom
+
+model_id = "poyum/test_discut"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+pipe = DiscoursePipeline(model_id=model_id, tokenizer=tokenizer)
+
+def predict(text):
+    return pipe(text)
+
+demo = gr.Interface(fn=predict, inputs="text", outputs="text")
+demo.launch()

pipeline.py

@@ -35,15 +35,15 @@ def write_sentences_to_format(sentences: list[str], filename: str):
 
 
 class DiscoursePipeline(Pipeline):
-    def __init__(self, model, tokenizer, config:dict, output_folder="./pipe_out",sat_model:str="sat-3l", **kwargs):
-        auto_model = AutoModelForTokenClassification.from_pretrained(model)
+    def __init__(self, model_id, tokenizer, output_folder="./pipe_out",sat_model:str="sat-3l", **kwargs):
+        auto_model = AutoModelForTokenClassification.from_pretrained(model_id)
         super().__init__(model=auto_model, tokenizer=tokenizer, **kwargs)
-        self.config = {"model_checkpoint": model, "sent_spliter":"sat","task":"seg","type":"tok","trace":False,"report_to":"none","sat_model":sat_model,"tok_config":{
+        self.config = {"model_checkpoint": model_id, "sent_spliter":"sat","task":"seg","type":"tok","trace":False,"report_to":"none","sat_model":sat_model,"tok_config":{
         "padding":"max_length",
         "truncation":True,
         "max_length": 512
     }}
-        self.model = model
+        self.model = model_id
         self.output_folder = output_folder
 
     def _sanitize_parameters(self, **kwargs):