Spaces:
Build error
Build error
# 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() |