Update app.py

app.py

@@ -1,9 +1,5 @@
-import time
 import gradio as gr
-from transformers import pipeline
-
-# Load the generation model
-generator = pipeline("text2text-generation", model="t5-small")
+from transformers import pipeline, set_seed
 
 # Define the questions
 questions = [
@@ -29,41 +25,57 @@ questions = [
     "What treatment plan do you recommend? Choose only from Options: (Scaling, Filling, Pulp therapy/RCT, Extraction, Medication, Referral) and nothing else"
 ]
 
-# Function to generate answers from context
-def generate_answers(context):
-    start_time = time.time()
-    answers = []
-    for question in questions:
-        # Generate the answer using the model
-        answer = generator(f"Context: {context}\nQuestion: {question}", max_length=50)[0]['generated_text']
-        answers.append(answer)
-    end_time = time.time()
-    generation_time = end_time - start_time
-    print(f"Time taken to generate answers: {generation_time:.2f} seconds")
-    return answers
+# Initialize answers
+answers = {q: '' for q in questions}
+
+# Define the default context
+default_context = (
+    "My child, Sarah, is studying in the 3rd grade. She is 8 years old and is a girl. "
+    "Sarah attends Greenwood Elementary School, located in Springfield. "
+    "Her guardians are myself, Emma Johnson, and her father, Michael Johnson.\n\n"
+    "Sarah's chief complaint regarding her oral health is that she has been experiencing some discomfort due to a cavity in her molar. "
+    "Other than that, she has no significant medication history. She has no relevant medical history to report. "
+    "Sarah does not take any medications regularly.\n\n"
+    "Her last dental visit was about a year ago, so this is her second visit. "
+    "Sarah has a habit of occasionally biting her nails, but she doesn't have any other habits like thumb sucking or tongue thrusting.\n\n"
+    "Sarah brushes her teeth twice daily. She does not experience bleeding gums. "
+    "Sarah has not experienced early childhood caries. Unfortunately, she does have tooth decay present in tooth numbers 14 and 15. "
+    "None of her teeth have been fractured. There is no pre-shedding mobility of teeth to report.\n\n"
+    "Sarah does not have malocclusion. She does not experience any pain, swelling, or abscess at the moment. "
+    "There are no other significant findings to note.\n\n"
+    "As for the treatment plan, I would recommend a filling for the decayed teeth."
+)
 
-# Gradio interface
-def main(context):
-    answers = generate_answers(context)
-    return answers
+# Load the GPT-2 model for text generation
+generator = pipeline('text-generation', model='gpt2')
+
+# Function to generate answers
+def generate_answers(context):
+    prompt = context + "\n\n" + "\n\n".join(questions)
+    set_seed(42)  # Ensure reproducibility
+    generated_text = generator(prompt, max_length=300, num_return_sequences=1)[0]['generated_text']
+    generated_answers = generated_text.split("\n\n")
+    return generated_answers
 
-# Gradio app layout
-with gr.Blocks() as demo:
-    context = gr.Textbox(lines=5, placeholder="Enter the context here", label="Context")
-    generate_button = gr.Button("Generate Answers")
-    output_table = gr.Dataframe(headers=["Question", "Answer"], row_count=20, col_count=2, visible=False)
-    save_button = gr.Button("Save Answers", visible=False)
+# Function to display the form
+def form_display():
+    return [gr.outputs.Textbox(label="Generated Answers", type="text")]
 
-    def generate_and_show_table(context):
-        answers = main(context)
-        data = [[question, answer] for question, answer in zip(questions, answers)]
-        generate_button.update(interactive=False)
-        output_table.update(data, visible=True)
-        save_button.update(visible=True)
-        return data
+# Function to submit the form
+def form_submit(context):
+    generated_answers = generate_answers(context)
+    return generated_answers
 
-    generate_button.click(generate_and_show_table, inputs=context, outputs=[output_table])
-    save_button.click(lambda: "Answers displayed in the table", inputs=None, outputs=gr.Textbox(label="Status"))
+# Create Gradio interface
+iface = gr.Interface(
+    fn=form_submit,
+    inputs=gr.inputs.Textbox(label="Context", type="text", default=default_context, lines=20),
+    outputs=form_display(),
+    title="Dental Health Record Form with GPT-2",
+    description="Fill out the context and click submit to generate answers to the questions.",
+    examples=[default_context],
+    allow_flagging=False
+)
 
-# Launch the app
-demo.launch()
+# Launch the interface
+iface.launch()