Update app.py

app.py

@@ -1,28 +1,18 @@
-import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
 
-# Load LLAMA model and tokenizer
-model_name = "sujra/insurance_Model"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
+import gradio as gr
+from transformers import pipeline
 
-# Define function for generating text
-#
-def generate_text(prompt):
-    pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200)
-    result = pipe(f"<s>[INST] {prompt} [/INST]")
-    generated_text = result[0]['generated_text']
-    return generated_text
+# Load your model from the Hugging Face model hub
+model_name = "sujra/insurance_Model"
+qa_pipeline = pipeline("question-answering", model=model_name, tokenizer=model_name)
 
-st.title("Insurance Response Generation")
+def get_answer(question):
+    answer = qa_pipeline(question=question, max_length=128, batch_size=4)['answer']  # Reduce batch size and set max_length
+    return answer
 
-prompt_input = st.text_input("Enter your prompt:")
+# Define the input and output components for the UI
+question_input = gr.inputs.Textbox(lines=2, label="Enter your question")
+output_text = gr.outputs.Textbox(label="Answer")
 
-if st.button("Generate Response"):
-    if prompt_input:
-        with st.spinner("Generating response..."):  # Display a spinner while generating response
-            response = generate_text(prompt_input)
-        st.write("Generated Response:")
-        st.write(response)
-    else:
-        st.write("Please enter a prompt.")
+# Create the UI
+gr.Interface(fn=get_answer, inputs=question_input, outputs=output_text, title="Insurance Question Answering").launch()