Spaces:
Sleeping
Sleeping
File size: 4,692 Bytes
f28768b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
from fastapi import FastAPI
from app import predict
import os
from huggingface_hub import login
from pydantic import BaseModel
import sys
from langchain.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
import os
import PyPDF2 as pdf
import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
import os
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings.sentence_transformer import (
SentenceTransformerEmbeddings,
)
from langchain_chroma import Chroma
from sentence_transformers import SentenceTransformer
from langchain_core.messages import AIMessage, HumanMessage
from fastapi import FastAPI, Request, UploadFile, File
os.environ['HF_HOME'] = '/hug/cache/'
os.environ['TRANSFORMERS_CACHE'] = '/blabla/cache/'
app = FastAPI()
app.recursion_limit = 10**4
def predict(message, db):
llm = ChatOpenAI(model_name="gpt-4o-mini", temperature=0)
template = """You are a general purpose chatbot. Be friendly and kind. Help people answer their questions. Use the context below to answer the questions
{context}
Question: {question}
Helpful Answer:"""
QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context", "question"],template=template,)
memory = ConversationBufferMemory(
memory_key="chat_history",
return_messages=True
)
retriever = db.as_retriever(k=3)
contextualize_q_system_prompt = """Given a chat history and the latest user question \
which might reference context in the chat history, formulate a standalone question \
which can be understood without the chat history. Do NOT answer the question, \
just reformulate it if needed and otherwise return it as is."""
contextualize_q_prompt = ChatPromptTemplate.from_messages(
[
("system", contextualize_q_system_prompt),
MessagesPlaceholder(variable_name="chat_history"),
("human", "{question}"),
]
)
contextualize_q_chain = contextualize_q_prompt | llm | StrOutputParser()
def contextualized_question(input: dict):
if input.get("chat_history"):
return contextualize_q_chain
else:
return input["question"]
rag_chain = (
RunnablePassthrough.assign(
context=contextualized_question | retriever
)
| QA_CHAIN_PROMPT
| llm
)
history = []
ai_msg = rag_chain.invoke({"question": message, "chat_history": history})
print(ai_msg)
bot_response = ai_msg.content.strip()
# Ensure history is correctly formatted as a list of tuples (user_message, bot_response)
history.append((HumanMessage(content=message), AIMessage(content=bot_response)))
docs = db.similarity_search(message,k=3)
extra = "\n" + "*"*100 + "\n"
additional_info = []
for d in docs:
citations = d.metadata["source"] + " pg." + str(d.metadata["page"])
additional_info = d.page_content
extra += citations + "\n" + additional_info + "\n" + "*"*100 + "\n"
# Return the bot's response and the updated history
return bot_response + extra
def upload_file(file_path):
loaders = []
print(file_path)
loaders.append(PyPDFLoader(file_path))
documents = []
for loader in loaders:
documents.extend(loader.load())
text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=16)
docs = text_splitter.split_documents(documents)
model = "thenlper/gte-large"
embedding_function = SentenceTransformerEmbeddings(model_name=model)
print(f"Model's maximum sequence length: {SentenceTransformer(model).max_seq_length}")
collection_name = "Autism"
persist_directory = "./chroma"
print(len(docs))
db = Chroma.from_documents(docs, embedding_function)
print("Done Processing, you can query")
return db
class Item(BaseModel):
code: str
@app.get("/")
async def root():
return {"Code Review Automation":"Version 1.0 'First Draft'"}
@app.post("/UploadFile/")
def predict(question: str, file: UploadFile = File(...)):
contents = file.file.read()
with open(file.filename, 'wb') as f:
f.write(contents)
db = upload_file(file.filename)
result = predict(question, db)
return {"answer":result}
|