Update app.py

app.py

@@ -51,13 +51,13 @@ def process_output(output):
 
 
 # Define the function for generating output based on input
-def generate_output(input_text):
+def generate_output(input_text,model):
     # Prompt for the instruction
     
     output = ""
     # Initialize the FastLanguageModel
     model, tokenizer = FastLanguageModel.from_pretrained(
-        model_name = "DataIntelligenceTeam/mistral-7b-instruct-v0.2-bnb-4bit",
+        model_name = model,
         max_seq_length = 2500,
         dtype = None,
         load_in_4bit = True,
@@ -108,20 +108,17 @@ examples = [
 
 model_options = ["DataIntelligenceTeam/NER-Phi-3-mini-4k-instruct"]
 #,"DataIntelligenceTeam/NER-gemma-7b-bnb-4bit","DataIntelligenceTeam/llama-3-8b-Instruct-bnb-4bit","DataIntelligenceTeam/mistral-7b-instruct-v0.2-bnb-4bit"]
-inputs = [
-    gr.inputs.Textbox(label="Input Text"),
-    #gr.inputs.Dropdown(label="Select the Fine-tuned Model", choices=["DataIntelligenceTeam/NER-Phi-3-mini-4k-instruct","DataIntelligenceTeam/llama-3-8b-Instruct-bnb-4bit"]),
-]
-
-outputs = [
-    gr.outputs.Textbox(label="Original Output Text"),
-    gr.outputs.Textbox(label="Formatted JSON")
-]
 
 # Create Gradio interface
 iface = gr.Interface(fn=generate_output, 
-                     inputs=inputs, 
-                     outputs= outputs,
+                     inputs=[
+                            gr.inputs.Textbox(label="Input Text"),
+                            gr.inputs.Dropdown(label="Select the Fine-tuned Model", choices=["DataIntelligenceTeam/NER-Phi-3-mini-4k-instruct","DataIntelligenceTeam/llama-3-8b-Instruct-bnb-4bit"])  
+                             ]
+                     outputs= [
+                            gr.outputs.Textbox(label="Original Output Text"),
+                            gr.outputs.Textbox(label="Formatted JSON")
+                            ],
                      #examples=examples,
                      title="Email Information Extraction", 
                      description="Extract key information from the provided email.")