Update app.py

app.py

@@ -7,7 +7,6 @@ import torch.optim as optim
 import onnx
 import onnxruntime
 import pandas as pd
-from sklearn.preprocessing import LabelEncoder
 from io import BytesIO
 
 # Define a simple neural network
@@ -35,11 +34,12 @@ class EmbeddingNN(nn.Module):
 
 def process_csv(csv_data):
     df = pd.read_csv(StringIO(csv_data))
-    library_encoder = LabelEncoder()
-    description_encoder = LabelEncoder()
-    df['library_encoded'] = library_encoder.fit_transform(df['library_name'])
-    df['description_encoded'] = description_encoder.fit_transform(df['description'])
-    return df, library_encoder, description_encoder
+    
+    # Replace LabelEncoder with custom encoding using pandas factorize
+    df['library_encoded'], library_classes = df['library_name'].factorize()
+    df['description_encoded'], description_classes = df['description'].factorize()
+
+    return df, library_classes, description_classes
 
 def train_and_export(df):
     model = EmbeddingNN(len(df['library_encoded'].unique()), len(df['description_encoded'].unique()))
@@ -59,13 +59,17 @@ def train_and_export(df):
     torch.onnx.export(model, torch.tensor([0], dtype=torch.long), buffer)
     return buffer
 
-def infer_from_onnx(model_buffer, library_name, library_encoder, description_encoder):
+def infer_from_onnx(model_buffer, library_name, library_classes, description_classes):
     byte_stream = BytesIO(model_buffer.getvalue())
     onnx_model = onnx.load(byte_stream)
     sess = onnxruntime.InferenceSession(byte_stream.getvalue())
-    encoded_library = library_encoder.transform([library_name])
-    outputs = sess.run(None, {sess.get_inputs()[0].name: encoded_library})
-    predicted_description = description_encoder.inverse_transform([outputs[0].argmax()])[0]
+
+    # Replace transform with custom encoding
+    encoded_library = torch.tensor([list(library_classes).index(library_name)], dtype=torch.long)
+    
+    outputs = sess.run(None, {sess.get_inputs()[0].name: encoded_library.numpy()})
+    predicted_description = description_classes[outputs[0].argmax()]
+
     return predicted_description
 
 # Streamlit UI
@@ -96,12 +100,11 @@ torch,PyTorch is an open-source machine learning library
 tensorflow,Open source software library for high performance numerical computations
 pandas,Data analysis and manipulation tool
 numpy,Library for numerical computations in Python
-scikit-learn,Machine learning library in Python
 """
 
 csv_data = st.text_area("Paste your CSV data here:", value=DEFAULT_CSV)
 if st.button('Convert CSV to ONNX Neural Net'):
-    df, library_encoder, description_encoder = process_csv(csv_data)
+    df, library_classes, description_classes = process_csv(csv_data)
     model_buffer = train_and_export(df)
     st.download_button(
         label="Download ONNX model",
@@ -114,5 +117,5 @@ if st.button('Convert CSV to ONNX Neural Net'):
 uploaded_model = st.file_uploader("Choose an ONNX model file for inference", type="onnx")
 library_name_to_infer = st.text_input("Enter a library name for inference:")
 if uploaded_model and library_name_to_infer:
-    prediction = infer_from_onnx(uploaded_model, library_name_to_infer, library_encoder, description_encoder)
+    prediction = infer_from_onnx(uploaded_model, library_name_to_infer, library_classes, description_classes)
     st.write(f"Predicted description: {prediction}")