kek

Autograder.py

@@ -0,0 +1,41 @@
+import streamlit as st
+from generator2 import response
+from retrieve import retrieve_molecule_index
+from PIL import Image
+
+st.title("Chem 210 Autograder")
+
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Use text_area for text input with a smaller height
+text_input = st.text_input("Indicate the molecule you want graded")
+
+# Use file_uploader for image input
+image_input = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
+
+if st.button("Submit"):
+    if text_input and image_input:
+        if text_input:
+            with st.chat_message("user"):
+                st.markdown(f"Does the following image indicate correct chemical structure of {text_input}?")
+            st.session_state.messages.append({"role": "user", "content": text_input})
+            index = int(retrieve_molecule_index(text_input)[0])
+            image_path = f"test/CID_{index}.png"
+            image2 = Image.open(image_path).convert('RGB')
+            
+        if image_input:
+            with st.chat_message("user"):
+                st.image(image_input, caption="User image", use_column_width=True)
+            st.session_state.messages.append({"role": "user", "content": "User uploaded an image."})
+            image = Image.open(image_input).convert('RGB')
+        
+        answer = response(image2, image)
+        with st.chat_message("AI"):
+            st.markdown(answer)
+        
+        st.session_state.messages.append({"role": "AI", "content": answer})

README.md

@@ -5,7 +5,7 @@ colorFrom: green
 colorTo: gray
 sdk: streamlit
 sdk_version: 1.36.0
-app_file: app.py
+app_file: Autograder.py
 pinned: false
 ---
 

config.ini

@@ -0,0 +1,3 @@
+[example]
+uri = https://in03-594505a48b55a63.api.gcp-us-west1.zillizcloud.com
+token = 4e989aded0471e56339d6d1ea894daac5435006f1b6bedd1fa5f1e38ac23aeade4e8b01797d397c2b2a07a292608fe7487b4c671

embeddings.py

@@ -0,0 +1,99 @@
+import pandas as pd
+import time
+import random
+from sentence_transformers import SentenceTransformer
+from pymilvus import connections, DataType, FieldSchema, CollectionSchema, Collection, utility
+import configparser
+from tqdm import tqdm
+
+# Initialize SentenceTransformer model for embeddings
+embedding_model = SentenceTransformer(model_name_or_path="bert-base-uncased")
+
+# Read molecule names from CSV
+csv_path = 'molecules-small.csv'
+df = pd.read_csv(csv_path)
+max_name_length = 256
+molecules = df['cmpdname'].tolist()
+
+for i, molecule in enumerate(molecules):
+    if len(molecule) > max_name_length:
+        molecules[i] = molecule[:max_name_length]
+        
+cids = df['cid'].tolist()
+
+# Encode embeddings for each molecule
+embeddings_list = []
+for molecule in tqdm(molecules, desc="Generating Embeddings"):
+    embeddings = embedding_model.encode(molecule)
+    embeddings_list.append(embeddings)
+
+cfp = configparser.RawConfigParser()
+cfp.read('config.ini')
+milvus_uri = cfp.get('example', 'uri')
+token = cfp.get('example', 'token')
+connections.connect("default",
+                    uri=milvus_uri,
+                    token=token)
+print(f"Connecting to DB: {milvus_uri}")
+# Define collection name and dimensionality of embeddings
+collection_name = 'molecule_embeddings'
+check_collection = utility.has_collection(collection_name)
+if check_collection:
+     drop_result = utility.drop_collection(collection_name)
+print("Success!")
+dim = 768  # Adjust based on the dimensionality of your embeddings
+
+# Define collection schema
+molecule_cid = FieldSchema(name="molecule_cid", dtype=DataType.INT64, description="cid", is_primary = True)
+molecule_name = FieldSchema(name="molecule_name", dtype=DataType.VARCHAR, max_length=256, description="name")
+molecule_embeddings = FieldSchema(name="molecule_embedding", dtype=DataType.FLOAT_VECTOR, dim=dim)
+schema = CollectionSchema(fields=[molecule_cid, molecule_name, molecule_embeddings],
+                      auto_id=False,
+                      description="my first collection!")
+
+print(f"Creating example collection: {collection_name}")
+collection = Collection(name=collection_name, schema=schema)
+print(f"Schema: {schema}")
+print("Success!")
+
+batch_size = 1000
+total_rt = 0
+start = 0
+print(f"Inserting {len(embeddings_list)} entities... ")
+for i in tqdm(range(0, len(embeddings_list), batch_size), desc="Inserting Embeddings"):
+    batch_embeddings = embeddings_list[i:i + batch_size]  
+    batch_molecules = molecules[i:i + batch_size]
+    batch_cids = cids[i:i + batch_size]
+    entities = [batch_cids, batch_molecules, batch_embeddings]
+    start += batch_size
+    t0 = time.time()
+    ins_resp = collection.insert(entities)
+    ins_rt = time.time() - t0
+    total_rt += ins_rt
+    
+
+print(f"Succeed in inserting {len(embeddings_list)} entities in {round(total_rt, 4)} seconds!")
+
+# Flush collection
+print("Flushing collection...")
+collection.flush()
+
+# Build index
+index_params = {"index_type": "AUTOINDEX", "metric_type": "L2", "params": {}}
+print("Building index...")
+collection.create_index(field_name='molecule_embedding', index_params=index_params)
+
+collection.load()
+    
+# Example search
+nq = 1
+search_params = {"metric_type": "L2"}
+topk = 5
+search_vec = [[random.random() for _ in range(dim)] for _ in range(nq)]
+print(f"Searching vector: {search_vec}")
+results = collection.search(search_vec, anns_field='molecule_embedding', param=search_params, limit=topk)
+print(f"Search results: {results}")
+
+# Disconnect from Milvus server
+connections.disconnect("default")
+print("Disconnected from Milvus server.")

generator2.py

@@ -0,0 +1,17 @@
+import torch
+from transformers import AutoProcessor, AutoModelForPreTraining
+
+def response(image2, image):
+
+    processor = AutoProcessor.from_pretrained("google/paligemma-3b-pt-224")
+    model = AutoModelForPreTraining.from_pretrained("google/paligemma-3b-pt-224")
+    
+    # Instruct the model to create a caption in Spanish
+    model_inputs = processor(text="check whether both molecules have the same chemical structure. if yes, output correct and if not, output incorrect", images= [image, image2], return_tensors="pt")
+    input_len = model_inputs["input_ids"].shape[-1]
+
+    with torch.inference_mode():
+        generation = model.generate(**model_inputs, max_new_tokens=100, do_sample=False)
+        generation = generation[0][input_len:]
+        decoded = processor.decode(generation, skip_special_tokens=True)
+    return decoded

pages/1_Manual_Image_Upload.py

@@ -0,0 +1,33 @@
+import streamlit as st
+from generator2 import response
+from PIL import Image
+
+st.title("Chem 210 Autograder - Manual Upload")
+
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Use file_uploader for image input
+image_input1 = st.file_uploader("Upload the solution molecule", type=["png", "jpg", "jpeg"])
+image_input2 = st.file_uploader("Upload the desired molecule to be graded", type=["png", "jpg", "jpeg"])
+
+if st.button("Submit"):
+    if image_input1 and image_input2:
+        # Open and convert images to RGB format
+        image1 = Image.open(image_input1).convert('RGB')
+        image2 = Image.open(image_input2).convert('RGB')
+
+        # Process images using your response function
+        answer = response(image1, image2)
+
+        # Display results in chat style
+        st.session_state.messages.append({"role": "user", "content": "User uploaded images."})
+        with st.chat_message("AI"):
+            st.markdown(answer)
+        st.session_state.messages.append({"role": "AI", "content": answer})
+    else:
+        st.warning("Please upload two images.")

pages/2_Retrieval_Demo.py

@@ -0,0 +1,28 @@
+import streamlit as st
+from generator2 import response
+from retrieve import retrieve_molecule_index
+from PIL import Image
+
+st.title("Retrieval System Demo")
+
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Use text_area for text input with a smaller height
+text_input = st.text_input("Indicate a desired molecule from our database")
+
+if st.button("Submit"):
+    if text_input:
+        with st.chat_message("user"):
+            st.markdown(f"Sending request to Milvus...")
+            index = int(retrieve_molecule_index(text_input)[0])
+            image_path = f"test/CID_{index}.png"
+            image = Image.open(image_path).convert('RGB')
+        
+        with st.chat_message("AI"):
+            st.write("Retrieved from our database:")
+            st.image(image, use_column_width=True)

pages/3_About.py

@@ -0,0 +1,50 @@
+import streamlit as st
+
+paragraph_1 = """
+The Autograder's underlying software design utilizes a RAG (Retrieval-Augmented Generation) pipeline.
+
+The context items consist of 3,717 2-d molecule images sourced from PubChem and saved into a 
+local directory(storing on a relational database is something to explore). Pubchem comes with a downloadable
+csv file that can easily be converted to a dataframe, which gives information on the molecule's cid and compound name.
+The embeddings are decoded through BERT and these embeddings are stored the vector database Milvus, where the molecule's
+cid served as the index.
+
+The retrieval system returns the cid of the molecule with the highest semantic score. The intuition behind this decision 
+choice was that molecules can have multiple names, so a simple keyword search would not be as versatile. Currently, it 
+has not gone through much testing but is able to produce accurate results sometimes when fed in a molecule compound
+synonym.
+
+The underlying VLM is Google/PaliGemma-3b and is currently not fine-tuned(hence the inaccurate results). The datasets that
+we have access to are too noisy and we simply do not have enough images to fine-tune a "lightweight" VLM such as 
+PaliGemma.
+
+"""
+paragraph_2 = """
+The Custom Image Upload Autograder's underlying software design employs a straightforward prompt engineering approach.
+
+The core model used is PaliGemma. This model was developed as a contingency measure to address potential reliability issues 
+that may arise with the retrieval system in the future.
+"""
+paragraph_3 = """
+The next step would be to fine-tune our underlying VLM. However, before this is feasible, we would need to gather
+a large corpus of images that is clean and organized well.
+
+To improve the retrieval system, there would need to be much more molecules uploaded and stored on Milvus. Mapping 
+these to a relational database would also be necessary is storing these on a local directory at this large of a scale
+is unfeasible.
+
+The Streamlit frontend could also be improved to make it more user-friendly and functional. This could include a
+better layout, more interactive features, and smooth integration with the backend for efficient data retrieval 
+and processing.
+"""
+
+st.title("About")
+st.markdown("""DISCLAIMER!
+            The underlying VLM is not finetuned, so the quality of the outputs are unreliable. 
+            """)
+st.header("About the Autograder")
+st.markdown(paragraph_1)
+st.header("About Manual Upload")
+st.markdown(paragraph_2)
+st.header("Moving Onwards")
+st.markdown(paragraph_3)

requirements.txt

@@ -0,0 +1,9 @@
+streamlit
+Pillow
+pandas
+sentence-transformers
+configparser
+pymilvus
+tqdm
+transformers
+torch

retrieve.py

@@ -0,0 +1,38 @@
+from pymilvus import Collection, DataType, FieldSchema, CollectionSchema, connections
+from sentence_transformers import SentenceTransformer
+import configparser
+
+
+def retrieve_molecule_index(molecule):
+    model = SentenceTransformer(model_name_or_path="bert-base-uncased")
+    search_vector = model.encode(molecule).reshape(1,-1)
+    cfp = configparser.RawConfigParser()
+    cfp.read('config.ini')
+    milvus_uri = cfp.get('example', 'uri')
+    token = cfp.get('example', 'token')
+    connections.connect("default",
+                        uri=milvus_uri,
+                        token=token)
+    print(f"Connecting to DB: {milvus_uri}")
+    collection_name = "molecule_embeddings"
+    dim = 768  # Adjust based on the dimensionality of your embeddings
+    
+    # Define collection schema
+    molecule_cid = FieldSchema(name="molecule_cid", dtype=DataType.INT64, description="cid", is_primary = True)
+    molecule_name = FieldSchema(name="molecule_name", dtype=DataType.VARCHAR, max_length=256, description="name")
+    molecule_embeddings = FieldSchema(name="molecule_embedding", dtype=DataType.FLOAT_VECTOR, dim=dim)
+    schema = CollectionSchema(fields=[molecule_cid, molecule_name, molecule_embeddings],
+                          auto_id=False,
+                          description="my first collection!")
+    
+    print(f"Creating example collection: {collection_name}")
+    collection = Collection(name=collection_name, schema=schema)
+    
+    search_params = {"metric_type": "IP"}
+    topk = 1
+    results = collection.search(search_vector, anns_field='molecule_embedding', param=search_params, limit=topk)
+    print(results)
+    # Disconnect from Milvus server
+    connections.disconnect("default")
+    print("Disconnected from Milvus server.")
+    return results[0].ids

setup.py

@@ -0,0 +1,11 @@
+import os
+
+index = 0
+for item in os.scandir("test"):
+    print(item)
+    index += 1
+    if index is 5:
+        break
+    
+    
+