Rag_chatbot / app.py
Michtiii's picture
Update app.py
98c04b2 verified
import gradio as gr
from langchain_text_splitters import CharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from transformers import pipeline
# -------------------------
# Load Data
# -------------------------
with open("data.txt", "r") as f:
raw_text = f.read()
# -------------------------
# Split Text
# -------------------------
text_splitter = CharacterTextSplitter(
chunk_size=300,
chunk_overlap=50
)
texts = text_splitter.split_text(raw_text)
# -------------------------
# Embeddings + Vector Store
# -------------------------
embeddings = HuggingFaceEmbeddings(
model_name="sentence-transformers/all-MiniLM-L6-v2"
)
vectorstore = FAISS.from_texts(texts, embeddings)
# -------------------------
# LLM (Stable HF Model)
# -------------------------
generator = pipeline(
"text-generation",
model="microsoft/phi-2",
max_new_tokens=120,
temperature=0.2,
do_sample=False
)
# -------------------------
# Chat Function (RAG + Memory)
# -------------------------
def rag_chat(message, history):
# Retrieve relevant docs
docs = vectorstore.similarity_search(message, k=3)
context = " ".join([doc.page_content for doc in docs])
# Build chat memory (last 3 messages)
history_text = ""
for user, bot in history[-3:]:
history_text += f"User: {user}\nBot: {bot}\n"
# Prompt
prompt = f"""
You are a professional AI assistant.
Rules:
- Answer ONLY from the context
- If answer not found, say "I don't know"
- Keep answer short and clear
Context:
{context}
Conversation:
{history_text}
Question:
{message}
Answer:
"""
# Generate response
result = generator(prompt)[0]["generated_text"]
# Clean output
answer = result.replace(prompt, "").strip()
answer = answer.split("\n")[0]
return answer
# -------------------------
# UI
# -------------------------
demo = gr.ChatInterface(
fn=rag_chat,
title="RAG Chatbot .txt base data Trained"
)
# -------------------------
# Run App
# -------------------------
if __name__ == "__main__":
demo.launch()