Update app.py

app.py

@@ -16,67 +16,62 @@ import gradio as gr
 from transformers import AutoModel, AutoTokenizer
 import torch
 
+# Load the model and tokenizer from Hugging Face
+tokenizer = AutoTokenizer.from_pretrained("Alibaba-NLP/gte-multilingual-base", trust_remote_code=True)
+model = AutoModel.from_pretrained("Alibaba-NLP/gte-multilingual-base", trust_remote_code=True)
+
 # Load the dataset containing PEC numbers and names
-def load_dataset(file_path='PEC_Numbers_and_Names.xlsx'):
-    if not os.path.exists(file_path):
+def load_dataset(file_path=None):
+    if file_path and os.path.exists(file_path):
+        df = pd.read_excel(file_path)
+    else:
         raise FileNotFoundError(f"File not found: {file_path}")
-    df = pd.read_excel(file_path)
     return df
 
-# Debugging function to get PEC number based on the name
-def get_pec_number(name, df):
-    print("Column names in DataFrame:", df.columns.tolist())  # Print the column names
-    print(f"Looking for Name: '{name}'")
+# Function to get the name based on the PEC number
+def get_name(pec_number, df):
+    df['PEC No.'] = df['PEC No.'].str.strip().str.upper()
+    pec_number = pec_number.strip().upper()
     
-    # Normalize the input and dataset
-    df['Name'] = df['Name'].str.strip().str.lower()
-    name = name.strip().str.lower()
-    
-    result = df[df['Name'] == name]
+    result = df[df['PEC No.'] == pec_number]
     
     if not result.empty:
-        print(f"Found PEC Number: {result.iloc[0]['PEC No.']}")
-        return result.iloc[0]['PEC No.']
+        return result.iloc[0]['Name']
     else:
-        print("Name not found.")
-        return "Name not found."
+        return "PEC Number not found."
 
-# Function to process the name using the Hugging Face model
-def process_with_model(name):
-    inputs = tokenizer(name, return_tensors="pt")
+# Function to process the PEC number using the Hugging Face model
+def process_with_model(pec_number):
+    inputs = tokenizer(pec_number, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
     return outputs.last_hidden_state.mean(dim=1).squeeze().tolist()
 
 # Combine both functions to create a prediction
-def predict(name, file):
+def predict(pec_number, file):
     try:
         # Load the dataset from the uploaded file if provided
         if file is not None:
             df = pd.read_excel(file.name)
         else:
-            df = load_dataset()
+            return "Please upload the PEC Numbers and Names file."
 
-        pec_number = get_pec_number(name, df)
-        model_output = process_with_model(name)
-        return f"PEC Number: {pec_number}\nModel Output: {model_output}"
+        name = get_name(pec_number, df)
+        model_output = process_with_model(pec_number)
+        return f"Name: {name}\nModel Output: {model_output}"
     except FileNotFoundError as e:
         return str(e)
 
-# Load the model and tokenizer from Hugging Face
-tokenizer = AutoTokenizer.from_pretrained("Alibaba-NLP/gte-multilingual-base", trust_remote_code=True)
-model = AutoModel.from_pretrained("Alibaba-NLP/gte-multilingual-base", trust_remote_code=True)
-
 # Build the Gradio interface with file upload option
 iface = gr.Interface(
     fn=predict,
     inputs=[
-        gr.Textbox(lines=1, placeholder="Enter Name..."),
-        gr.File(label="Upload PEC Numbers and Names file (optional)")
+        gr.Textbox(lines=1, placeholder="Enter PEC Number..."),
+        gr.File(label="Upload PEC Numbers and Names file")
     ],
     outputs="text",
-    title="Name to PEC Number Lookup with Model Integration",
-    description="Enter a name to retrieve the corresponding PEC number and process it with a Hugging Face model. Optionally, upload the Excel file if not found."
+    title="PEC Number to Name Lookup",
+    description="Enter a PEC number to retrieve the corresponding name and process it with a Hugging Face model. Please upload the Excel file."
 )
 
 # Run the Gradio interface