Update app.py

app.py

@@ -5,6 +5,7 @@ import traceback
 import pandas as pd
 import torch
 import gradio as gr
+import gc
 from transformers import (
     logging,
     AutoProcessor,
@@ -106,6 +107,27 @@ def generate_and_export():
             "outputs/inference.tsv",
             "outputs/eval.csv"
         )
+        # Add this to your generate_and_export function
+
+        # ... existing code ...
+        
+        # Better memory management
+        for i in range(1, 21):
+            doc = to_soap("Generate a realistic, concise doctor's progress note...")
+            docs.append(doc)
+            gts.append(to_soap(doc))
+            
+            # More aggressive cleanup every 3 iterations
+            if i % 3 == 0:
+                torch.cuda.empty_cache()
+                gc.collect()
+        
+        # Clean up model after use
+        del model
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+
 
     except Exception as e:
         traceback.print_exc()