Update app.py

app.py

@@ -1,46 +1,22 @@
-import gradio as gr
 import pandas as pd
-from sentence_transformers import SentenceTransformer
-import chromadb
-from transformers import GPT2LMHeadModel, GPT2Tokenizer, TrainingArguments, Trainer
-from datasets import Dataset
-from torch.cuda.amp import autocast
-import torch
-
-def load_and_process_data(file_info):
-    if file_info is None:
-        return None, "Please upload a CSV file."
-    data = pd.read_csv(file_info, nrows=500, on_bad_lines='skip')
-    data['message'] = data['message'].apply(lambda x: x.strip() if type(x) == str else '')
-    return data, "Data loaded and processed successfully."
-
-def answer_question(model, tokenizer, question):
-    model.eval()
-    inputs = tokenizer.encode(question, return_tensors='pt')
-    outputs = model.generate(inputs, max_length=100, num_return_sequences=1)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-tokenizer = GPT2Tokenizer.from_pretrained('distilgpt2')
-tokenizer.pad_token = tokenizer.eos_token
-model = GPT2LMHeadModel.from_pretrained('distilgpt2')
-
-def setup_gradio_interface():
-    with gr.Blocks() as demo:
-        gr.Markdown("### Upload emails.csv and Ask a Question")
-        with gr.Row():
-            file_input = gr.File(label="Upload your emails.csv")
-            submit_file = gr.Button("Load & Process File")
-        data_output = gr.Textbox(label="File Processing Output")
-        question_input = gr.Textbox(label="Enter your question:")
-        answer_output = gr.Textbox(label="Model Answer:")
-        submit_question = gr.Button("Get Answer")
-
-        submit_file.click(fn=load_and_process_data, inputs=file_input, outputs=[data_output])
-        submit_question.click(fn=answer_question, inputs=[model, tokenizer, question_input], outputs=answer_output)
-
-    return demo
-
-demo = setup_gradio_interface()
-
-if __name__ == "__main__":
-    demo.launch()
+import gradio as gr
+from transformers import pipeline
+
+emails_df = pd.read_csv('emails.csv', nrows=500, on_bad_lines='skip')
+context = " ".join(emails_df['message'].apply(lambda x: x.strip() if isinstance(x, str) else ''))
+
+qa_pipeline = pipeline("question-answering")
+
+def answer_query(question):
+    try:
+        result = qa_pipeline(question=question, context=context)
+        return result['answer']
+    except Exception as e:
+        return str(e)  
+
+iface = gr.Interface(
+    fn=answer_query,
+    inputs=gr.inputs.Textbox(label="Enter your question:"),
+    outputs=gr.outputs.Textbox(label="Answer:")
+)
+iface.launch()