|
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 |
|
|
|
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY") |
|
PINECONE_INDEX = "arolchatbot" |
|
pc = PC(api_key=PINECONE_API_KEY) |
|
index = pc.Index(PINECONE_INDEX) |
|
|
|
|
|
|
|
embeddings = HuggingFaceEmbeddings(model_name="thenlper/gte-large") |
|
vector_store = PineconeVectorStore(index, embeddings, "text") |
|
|
|
|
|
model_name = "Meldashti/chatbot" |
|
base_model = AutoModelForCausalLM.from_pretrained("unsloth/Llama-3.2-3B") |
|
tokenizer = AutoTokenizer.from_pretrained("unsloth/Llama-3.2-3B") |
|
|
|
|
|
model = PeftModel.from_pretrained(base_model, model_name) |
|
model = model.merge_and_unload() |
|
|
|
|
|
generator = pipeline( |
|
"text-generation", |
|
model=model, |
|
tokenizer=tokenizer, |
|
max_new_tokens=150 |
|
) |
|
|
|
|
|
llm = HuggingFacePipeline(pipeline=generator) |
|
|
|
retriever = vector_store.as_retriever(search_kwargs={"k": 5}) |
|
|
|
|
|
rag_chain = RetrievalQA.from_chain_type( |
|
llm=llm, |
|
chain_type="stuff", |
|
retriever=retriever |
|
) |
|
|
|
|
|
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)}" |
|
|
|
|
|
demo = gr.ChatInterface(chat, type="messages", autofocus=False) |
|
|
|
|
|
if __name__ == "__main__": |
|
demo.launch(debug=True) |