File size: 1,998 Bytes
b560279 58bf06b b560279 58bf06b 0e7f45b a3272dc 58bf06b 73768fb 58bf06b a2697fb 58bf06b b560279 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import gradio as gr
from sentence_transformers import SentenceTransformer
from transformers import pipeline
import faiss
import os
# Load models
embedder = SentenceTransformer("all-MiniLM-L6-v2")
#qa_pipeline = pipeline("text-generation", model="gpt2") # Replace with a better QA model if desired
qa_pipeline = pipeline("text-generation", model="tiiuae/falcon-rw-1b")
# In-memory vector index and docs
index = None
docs = []
def process_docs(files):
global index, docs
docs = []
# Read files and compute embeddings
texts = []
for file in files:
with open(file.name, 'r', encoding='utf-8') as f:
content = f.read()
texts.append(content)
docs.append(content)
embeddings = embedder.encode(texts, convert_to_tensor=False)
dim = embeddings[0].shape[0]
# Create FAISS index
index = faiss.IndexFlatL2(dim)
index.add(embeddings)
return "Documents uploaded and indexed!"
def chat(query):
if not index or not docs:
return "Please upload documents first."
# Embed query and search
q_embedding = embedder.encode([query])
D, I = index.search(q_embedding, k=3)
# Get top documents
context = "\n".join([docs[i] for i in I[0]])
prompt = f"Context: {context}\nQuestion: {query}\nAnswer:"
response = qa_pipeline(prompt, max_new_tokens=100, do_sample=False)
return response[0]['generated_text'].split("Answer:")[-1].strip()
# Gradio UI
with gr.Blocks() as demo:
with gr.Row():
file_upload = gr.File(file_types=[".txt"], file_count="multiple")
upload_btn = gr.Button("Upload Documents")
upload_output = gr.Textbox(label="Status")
with gr.Row():
question = gr.Textbox(label="Ask a Question")
answer = gr.Textbox(label="Answer")
ask_btn = gr.Button("Ask")
upload_btn.click(fn=process_docs, inputs=[file_upload], outputs=[upload_output])
ask_btn.click(fn=chat, inputs=[question], outputs=[answer])
demo.launch()
|