Create app.py

app.py

@@ -0,0 +1,57 @@
+import random
+from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline
+from datasets import load_dataset
+import gradio as gr
+
+# Load the dataset with streaming
+dataset = load_dataset("TheBritishLibrary/blbooks", split="train", streaming=True)
+
+# Convert streaming dataset to an iterable
+dataset_iter = iter(dataset)
+
+# Load tokenizer and model
+model_name = "max-long/textile_machines_3_oct"  # Replace with your model's name
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForTokenClassification.from_pretrained(model_name)
+
+# Initialize NER pipeline
+ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
+
+def get_random_snippet(stream_iter, tokenizer, max_tokens=350, max_attempts=1000):
+    for _ in range(max_attempts):
+        try:
+            sample = next(stream_iter)['text']
+            tokens = tokenizer.tokenize(sample)
+            if len(tokens) <= max_tokens:
+                return sample
+        except StopIteration:
+            break
+    return "No suitable snippet found."
+
+def extract_textile_machinery_entities(text):
+    ner_results = ner_pipeline(text)
+    textile_entities = [ent for ent in ner_results if ent['entity_group'] == 'TEXTILE_MACHINERY']
+    return textile_entities
+
+def analyze_text():
+    snippet = get_random_snippet(dataset_iter, tokenizer)
+    entities = extract_textile_machinery_entities(snippet)
+    
+    # Highlight entities in the text
+    for ent in sorted(entities, key=lambda x: x['start'], reverse=True):
+        snippet = snippet[:ent['start']] + f"**{snippet['start']:ent['end']}**" + snippet[ent['end']:]
+    
+    return snippet, entities
+
+# Build Gradio interface
+with gr.Blocks() as demo_interface:
+    gr.Markdown("# Textile Machinery Entity Recognition Demo")
+    gr.Markdown("Click the button below to analyze a random text snippet.")
+    with gr.Row():
+        analyze_button = gr.Button("Analyze Random Snippet")
+    output_text = gr.Markdown()
+    output_entities = gr.JSON()
+    
+    analyze_button.click(fn=analyze_text, outputs=[output_text, output_entities])
+
+demo_interface.launch()