File size: 2,942 Bytes
06e0c89 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
class TextAnalyzer:
def __init__(self):
# Load the medical text analysis model
try:
model_name = "medicalai/ClinicalBERT"
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(self.device)
# NER pipeline for medical entities
self.ner_pipeline = pipeline(
"ner", model="samrawal/bert-base-uncased_medical-ner"
)
print(f"Text model loaded on {self.device}")
except Exception as e:
print(f"Error loading text model: {e}")
# Fallback to dummy functionality
self.model = None
self.tokenizer = None
self.ner_pipeline = None
def analyze(self, text):
"""Analyze medical report text and extract key insights"""
if text.strip() == "":
return {"Insights": "No text provided"}
if self.model is None or self.tokenizer is None:
# Dummy analysis
return {
"Entities": ["fever", "cough"],
"Sentiment": "Concerning",
"Key findings": "Patient shows symptoms of respiratory illness",
}
try:
# Extract medical entities
if self.ner_pipeline:
entities = self.ner_pipeline(text)
unique_entities = list(set([entity["word"] for entity in entities]))
else:
unique_entities = []
# Simple text classification (in real app, would be more sophisticated)
inputs = self.tokenizer(
text, return_tensors="pt", padding=True, truncation=True
).to(self.device)
with torch.no_grad():
outputs = self.model(**inputs)
# This is a placeholder - in reality would depend on the actual model output
sentiment = (
"Concerning" if torch.sigmoid(outputs.logits).item() > 0.5 else "Normal"
)
# Generate key findings (simplified)
key_findings = f"Report indicates {'abnormal' if sentiment == 'Concerning' else 'normal'} findings"
if unique_entities:
key_findings += f" with mentions of {', '.join(unique_entities[:5])}"
return {
"Entities": unique_entities[:10],
"Sentiment": sentiment,
"Key findings": key_findings,
}
except Exception as e:
print(f"Error during text analysis: {e}")
return {"Error": "Could not analyze text"}
|