tryagain / app.py
random2222's picture
Update app.py
eb90789 verified
raw
history blame
3.12 kB
# app.py (CPU-optimized)
import gradio as gr
import os
import torch
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from transformers import AutoModelForCausalLM, AutoTokenizer
# Configuration
DOCS_DIR = "business_docs"
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
MODEL_NAME = "microsoft/phi-2"
def initialize_system():
# Document processing
if not os.path.exists(DOCS_DIR):
raise FileNotFoundError(f"Missing {DOCS_DIR} folder")
pdf_files = [os.path.join(DOCS_DIR, f)
for f in os.listdir(DOCS_DIR)
if f.endswith(".pdf")]
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=500, # Smaller chunks for CPU
chunk_overlap=50
)
texts = []
for pdf in pdf_files:
loader = PyPDFLoader(pdf)
pages = loader.load_and_split(text_splitter)
texts.extend(pages)
# Create embeddings
embeddings = HuggingFaceEmbeddings(
model_name=EMBEDDING_MODEL,
model_kwargs={'device': 'cpu'}
)
# Vector store
vector_store = FAISS.from_documents(texts, embeddings)
# Load model without quantization
tokenizer = AutoTokenizer.from_pretrained(
MODEL_NAME,
trust_remote_code=True,
padding_side="left"
)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
trust_remote_code=True,
torch_dtype=torch.float16,
device_map="cpu" # Force CPU
)
return vector_store, model, tokenizer
try:
vector_store, model, tokenizer = initialize_system()
print("βœ… System initialized successfully")
except Exception as e:
print(f"❌ Initialization failed: {str(e)}")
raise
def generate_response(query):
try:
docs = vector_store.similarity_search(query, k=1) # Less context
context = "\n".join([d.page_content for d in docs])
prompt = f"""<|system|>
Answer using: {context}
- Max 1 sentence
- If unsure: "I'll check with the team"</s>
<|user|>{query}</s>
<|assistant|>"""
inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
outputs = model.generate(
**inputs,
max_new_tokens=100,
temperature=0.1
)
return tokenizer.decode(outputs[0], skip_special_tokens=True).split("<|assistant|>")[-1].strip()
except Exception as e:
return "Please try again later."
# Simplified interface
with gr.Blocks() as demo:
gr.Markdown("# Customer Service Chatbot")
chatbot = gr.Chatbot()
msg = gr.Textbox(label="Your question")
clear = gr.ClearButton([msg, chatbot])
def respond(message, history):
response = generate_response(message)
history.append((message, response))
return "", history
msg.submit(respond, [msg, chatbot], [msg, chatbot])
demo.launch()