Update app.py

app.py

@@ -16,16 +16,19 @@ import difflib
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # OCR Correction Model
-ocr_model_name = "PleIAs/OCRonos"
-ocr_llm = LLM(ocr_model_name, max_model_len=8128)
+ocr_model_name = "PleIAs/OCRonos-Vintage"
 
-# Editorial Segmentation Model
-editorial_model = "PleIAs/Segmentext"
-token_classifier = pipeline(
-    "token-classification", model=editorial_model, aggregation_strategy="simple", device=device
-)
+import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+
+# Load pre-trained model and tokenizer
+model_name = "PleIAs/OCRonos-Vintage"
+model = GPT2LMHeadModel.from_pretrained(model_name)
+tokenizer = GPT2Tokenizer.from_pretrained(model_name)
 
-tokenizer = AutoTokenizer.from_pretrained(editorial_model, model_max_length=512)
+# Set the device to GPU if available, otherwise use CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
 # CSS for formatting
 css = """
@@ -163,30 +166,25 @@ def split_text(text, max_tokens=500):
 
     return chunks
 
-def transform_chunks(marianne_segmentation):
-    marianne_segmentation = pd.DataFrame(marianne_segmentation)
-    marianne_segmentation = marianne_segmentation[marianne_segmentation['entity_group'] != 'separator']
-    marianne_segmentation['word'] = marianne_segmentation['word'].astype(str).str.replace('¶', '\n', regex=False)
-    marianne_segmentation['word'] = marianne_segmentation['word'].astype(str).apply(preprocess_text)
-    marianne_segmentation = marianne_segmentation[marianne_segmentation['word'].notna() & (marianne_segmentation['word'] != '') & (marianne_segmentation['word'] != ' ')]
 
-    html_output = []
-    for _, row in marianne_segmentation.iterrows():
-        entity_group = row['entity_group']
-        result_entity = "[" + entity_group.capitalize() + "]"
-        word = row['word']
-        
-        if entity_group == 'title':
-            html_output.append(f'<div class="manuscript"><div class="annotation">{result_entity}</div><div class="content title-content"><h2>{word}</h2></div></div>')
-        elif entity_group == 'bibliography':
-            html_output.append(f'<div class="manuscript"><div class="annotation">{result_entity}</div><div class="content bibliography-content">{word}</div></div>')
-        elif entity_group == 'paratext':
-            html_output.append(f'<div class="manuscript"><div class="annotation">{result_entity}</div><div class="content paratext-content">{word}</div></div>')
-        else:
-            html_output.append(f'<div class="manuscript"><div class="annotation">{result_entity}</div><div class="content">{word}</div></div>')
+# Function to generate text
+def ocr_correction(prompt, max_new_tokens=600):
+
+    prompt = f"""### Text ###\n{prompt}\n\n\n### Correction ###\n"""
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
+
+    # Generate text
+    output = model.generate(input_ids,
+                            max_new_tokens=max_new_tokens,
+                            pad_token_id=tokenizer.eos_token_id,
+                            top_k=50)
 
-    final_html = '\n'.join(html_output)
-    return final_html
+    # Decode and return the generated text
+    result = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    result = result.split("### Correction ###\n")[1]
+
+    return result
 
 # OCR Correction Class
 class OCRCorrector:
@@ -195,53 +193,25 @@ class OCRCorrector:
 
     def correct(self, user_message):
         sampling_params = SamplingParams(temperature=0.9, top_p=0.95, max_tokens=4000, presence_penalty=0, stop=["#END#"])
-        detailed_prompt = f"### TEXT ###\n{user_message}\n\n### CORRECTION ###\n"
-        prompts = [detailed_prompt]
-        outputs = ocr_llm.generate(prompts, sampling_params, use_tqdm=False)
-        generated_text = outputs[0].outputs[0].text
+        detailed_prompt = f"### Text ###\n{user_message}\n\n### Correction ###\n"
+        generated_text = ocr_correction(detailed_prompt)
         html_diff = generate_html_diff(user_message, generated_text)
         return generated_text, html_diff
 
-# Editorial Segmentation Class
-class EditorialSegmenter:
-    def segment(self, text):
-        editorial_text = re.sub("\n", " ¶ ", text)
-        num_tokens = len(tokenizer.tokenize(editorial_text))
-        
-        if num_tokens > 500:
-            batch_prompts = split_text(editorial_text, max_tokens=500)
-        else:
-            batch_prompts = [editorial_text]
-    
-        out = token_classifier(batch_prompts)
-        classified_list = []
-        for classification in out:
-            df = pd.DataFrame(classification)
-            classified_list.append(df)
-    
-        classified_list = pd.concat(classified_list)
-        out = transform_chunks(classified_list)
-        return out
-
 # Combined Processing Class
 class TextProcessor:
     def __init__(self):
         self.ocr_corrector = OCRCorrector()
-        self.editorial_segmenter = EditorialSegmenter()
 
     @spaces.GPU(duration=120)
     def process(self, user_message):
-        # Step 1: OCR Correction
+        #OCR Correction
         corrected_text, html_diff = self.ocr_corrector.correct(user_message)
         
-        # Step 2: Editorial Segmentation
-        segmented_text = self.editorial_segmenter.segment(corrected_text)
-        
         # Combine results
         ocr_result = f'<h2 style="text-align:center">OCR Correction</h2>\n<div class="generation">{html_diff}</div>'
-        editorial_result = f'<h2 style="text-align:center">Editorial Segmentation</h2>\n<div class="generation">{segmented_text}</div>'
         
-        final_output = f"{css}{ocr_result}<br><br>{editorial_result}"
+        final_output = f"{css}{ocr_result}"
         return final_output
 
 # Create the TextProcessor instance
@@ -249,7 +219,7 @@ text_processor = TextProcessor()
 
 # Define the Gradio interface
 with gr.Blocks(theme='JohnSmith9982/small_and_pretty') as demo:
-    gr.HTML("""<h1 style="text-align:center">PleIAs Editor</h1>""")
+    gr.HTML("""<h1 style="text-align:center">Vintage OCR corrector</h1>""")
     text_input = gr.Textbox(label="Your (bad?) text", type="text", lines=5)
     process_button = gr.Button("Process Text")
     text_output = gr.HTML(label="Processed text")