create main.py

main.py

@@ -0,0 +1,86 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("aaditya/Llama3-OpenBioLLM-8B")
+model = AutoModelForCausalLM.from_pretrained(
+    "aaditya/Llama3-OpenBioLLM-8B",
+    torch_dtype=torch.bfloat16,  # Use bfloat16 for better memory efficiency
+    device_map="auto"  # Automatically handle device placement
+)
+
+# Define the system prompt
+SYSTEM_PROMPT = """You are an advanced medical documentation assistant analyzing clinical documentation. For each case, provide:
+Documentation Quality Score (1-10):
+- Assess completeness, clarity, and adherence to clinical standards
+Diagnostic Confidence (0-100%):
+- Evaluate certainty of current diagnosis
+- Consider supporting evidence
+Key Findings:
+- List significant clinical observations
+- Note critical vital signs
+- Highlight abnormal results
+Missing Information:
+- Identify crucial missing data
+- Note incomplete documentation areas
+Recommended Actions:
+- Suggest immediate clinical steps
+- Propose management changes
+Additional Tests:
+- Recommend relevant diagnostics
+- Suggest appropriate imaging
+- Propose lab work
+Safety Concerns:
+- Flag potential drug interactions
+- Highlight clinical red flags
+- Note urgent safety issues"""
+
+def analyze_medical_doc(medical_doc):
+    # Prepare the prompt
+    prompt = f"{SYSTEM_PROMPT}\n\nPlease analyze this medical documentation:\n\n{medical_doc}\n\nAnalysis:"
+    
+    # Tokenize input
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
+    inputs = inputs.to(model.device)  # Move to same device as model
+    
+    # Generate response
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=512,
+        temperature=0.7,
+        top_p=0.9,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
+    )
+    
+    # Decode and clean up the response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract only the generated part (remove the input prompt)
+    response = response[len(prompt):]
+    
+    return response.strip()
+
+# Example usage
+sample_doc = """
+Patient: 45-year-old male
+Chief Complaint: Chest pain and shortness of breath
+Vitals:
+- BP: 145/90
+- HR: 98
+- RR: 20
+- O2 Sat: 95% on RA
+
+History:
+- Onset: 2 hours ago
+- Character: Sharp, radiating to left arm
+- Previous MI: No
+- HTN: Yes, on lisinopril
+
+Current Medications:
+- Lisinopril 10mg daily
+"""
+
+# Get the analysis
+analysis = analyze_medical_doc(sample_doc)
+print(analysis)