Update app.py

app.py

@@ -5,15 +5,16 @@ import torchvision
 from transformers import pipeline
 auth_token = os.environ.get("HUGGING_FACE_HUB_TOKEN")
 
-# Initialize the pipeline
-pipe = pipeline(
-    "text-generation",
-    model="Jyotiyadav/mistral_7B_NER",
-    torch_dtype=torch.bfloat16,
-    device_map="auto")
-
 # Function to generate output based on input
-def generate_output(input_text):
+def generate_output(input_text, max_new_tokens, temperature, top_k, top_p, model):
+    # Initialize the pipeline
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        torch_dtype=torch.bfloat16,
+        device_map="auto"
+    )
+    
     # Prompt for extracting information
     prompt = f'''
     Your task is to extract the information corresponding to the provided labels from the below given email.
@@ -52,16 +53,17 @@ def generate_output(input_text):
     result = pipe(
         f"<s>[INST] {prompt} [/INST]",
         do_sample=True,
-        max_new_tokens=32000,
-        temperature=0.1,
-        top_k=0,
-        top_p=0,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_k=top_k,
+        top_p=top_p,
         num_return_sequences=1,
     )
 
     # Return the generated text
     return result[0]['generated_text']
 
+
 examples = [
 '''    
 COTIZACION FLETE MARITIMO OC 4500325343 Buongiorno ;
@@ -89,13 +91,26 @@ Del. Alvaro Obregon ;
 
 
 
+# Create Gradio inputs
+inputs = [
+    gr.inputs.Textbox(label="Input Text"),
+    gr.inputs.Number(label="Max New Tokens", default=32000),
+    gr.inputs.Slider(label="Temperature", minimum=0.0, maximum=1.0, default=0.1, step=0.01),
+    gr.inputs.Number(label="Top K", default=0),
+    gr.inputs.Number(label="Top P", default=0),
+    gr.inputs.Textbox(label="Model", default="Jyotiyadav/mistral_7B_NER")
+]
 
 # Create a Gradio interface
-iface = gr.Interface(fn=generate_output, 
-                     inputs=["text"],
-                     outputs="text", 
-                     title="Information Extraction with Mistral-7B",
-                     examples=examples,
-                     debug=True,
-                     description="Generate a Information Extraction with OpenLLM.")
-iface.launch()
+iface = gr.Interface(
+    fn=generate_output, 
+    inputs=inputs,
+    outputs="text",
+    #examples=examples,
+    title="Information Extraction with Mistral-7B",
+    description="Generate Information Extraction with OpenLLM.",
+    debug=True
+)
+
+# Launch the interface
+iface.launch()