Engineer / app.py
atifsial123's picture
Update app.py
6feb7ec verified
raw
history blame
2.51 kB
import os
import subprocess
# Function to install a package if it is not already installed
def install(package):
subprocess.check_call([os.sys.executable, "-m", "pip", "install", package])
# Ensure the necessary packages are installed
install("transformers")
install("torch")
install("pandas")
install("gradio")
import pandas as pd
import gradio as gr
from transformers import AutoModel, AutoTokenizer
import torch
# 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):
raise FileNotFoundError(f"File not found: {file_path}")
df = pd.read_excel(file_path)
return df
# 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)
# Define the function to get the name based on the PEC number
def get_name(pec_number, df):
result = df[df['PEC No.'] == pec_number]
if not result.empty:
return result.iloc[0]['Name']
else:
return "PEC Number not found."
# 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(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()
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)
# Build the Gradio interface with file upload option
iface = gr.Interface(
fn=predict,
inputs=[
gr.Textbox(lines=1, placeholder="Enter PEC Number..."),
gr.File(label="Upload PEC Numbers and Names file (optional)")
],
outputs="text",
title="PEC Number Lookup with Model Integration",
description="Enter a PEC number to retrieve the corresponding name and process it with a Hugging Face model. Optionally, upload the Excel file if not found."
)
# Run the Gradio interface
if __name__ == "__main__":
iface.launch()