mohammadhakimi's picture
Update app.py
c9baad3 verified
raw
history blame
2.28 kB
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from peft import PeftModel
from langchain_community.vectorstores import FAISS
from langchain.text_splitter import CharacterTextSplitter
from langchain.docstore.document import Document
from langchain.chains import RetrievalQA
from langchain_huggingface import HuggingFacePipeline
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from sentence_transformers import SentenceTransformer
import os
import numpy as np
from langchain.schema import Document
from pinecone import Pinecone as PC
from langchain_pinecone import PineconeVectorStore
# Initialize Pinecone
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
PINECONE_INDEX = "arolchatbot" # e.g., "us-west1-gcp-free"
pc = PC(api_key=PINECONE_API_KEY)
index = pc.Index(PINECONE_INDEX)
# Connect to Pinecone
embeddings = HuggingFaceEmbeddings(model_name="thenlper/gte-large")
vector_store = PineconeVectorStore(index, embeddings, "text")
# Model and Tokenizer
model_name = "Meldashti/chatbot"
base_model = AutoModelForCausalLM.from_pretrained("unsloth/Llama-3.2-3B")
tokenizer = AutoTokenizer.from_pretrained("unsloth/Llama-3.2-3B")
# Merge PEFT weights with base model
model = PeftModel.from_pretrained(base_model, model_name)
model = model.merge_and_unload()
# Simplified pipeline with minimal parameters
generator = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_new_tokens=150
)
# LLM wrapper
llm = HuggingFacePipeline(pipeline=generator)
retriever = vector_store.as_retriever(search_kwargs={"k": 5})
# Retrieval QA Chain
rag_chain = RetrievalQA.from_chain_type(
llm=llm,
chain_type="stuff",
retriever=retriever
)
# Chat function with extensive logging
def chat(message, history):
print(f"Received message: {message}")
try:
response = rag_chain.invoke({"query": message})
print(response)
return str(response['result'].split("Helpful Answer: ")[1])
except Exception as e:
print(f"Error generating response: {type(e)}, {e}")
return f"An error occurred: {str(e)}"
# Gradio interface
demo = gr.ChatInterface(chat, type="messages", autofocus=False)
# Launch
if __name__ == "__main__":
demo.launch(debug=True)