Update app.py

app.py

@@ -10,6 +10,7 @@ import base64
 from docx import Document
 from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
 from io import BytesIO
+import re
 
 openai.api_key = "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"
 
@@ -64,7 +65,8 @@ def show_predicted_caption(image, top_k=1):
     matches = predict_caption(
         image, model, text_embeddings, testing_df["caption"]
     )[:top_k]
-    return matches[0].split('(')[0].strip()  # Remove the reference string from the caption
+    cleaned_matches = [re.sub(r'\s\(ROCO_\d+\)', '', match) for match in matches]  # Add this line to clean the matches
+    return cleaned_matches  # Return the cleaned_matches instead of matches
 
 def generate_radiology_report(prompt):
     response = openai.Completion.create(
@@ -75,7 +77,11 @@ def generate_radiology_report(prompt):
         stop=None,
         temperature=1,
     )
-    return response.choices[0].text.strip()
+    report = response.choices[0].text.strip()
+    # Remove reference string from the report
+    report = re.sub(r'\(ROCO_\d+\)', '', report).strip()
+    return report
+
 
 def save_as_docx(text, filename):
     document = Document()