Spaces:
Sleeping
Sleeping
File size: 4,089 Bytes
93461a2 e25d199 93461a2 cc6c7e9 93461a2 396d74c 93461a2 aba63c5 93461a2 506fac0 f753eb9 6d78ad1 93461a2 | 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 | import os
from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain.text_splitter import RecursiveCharacterTextSplitter
import streamlit as st
from together import Together
from langchain.llms.base import LLM
from typing import Any, List, Optional
from Components.para_utility import load_pdfs_from_file
from pydantic import PrivateAttr
embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
class TogetherLLM(LLM):
model_name: str = "mistralai/Mistral-7B-Instruct-v0.2"
temperature: float = 0
max_tokens: int = 256
together_api_key: str = os.getenv("Together_API")
# Define private attribute
client: Any = PrivateAttr()
def __init__(self, **kwargs):
super().__init__(**kwargs)
# Together("api_key")=self.together_api_key
self.client = Together(api_key=self.together_api_key)
def _call(self, prompt: str, **kwargs: Any) -> str:
response = self.client.chat.completions.create(
model=self.model_name,
messages=[{"role": "user", "content": prompt}],
temperature=self.temperature,
max_tokens=self.max_tokens,
)
return response.choices[0].message.content
@property
def _llm_type(self) -> str:
return "together_llm"
def split_docs(documents, chunk_size=500, chunk_overlap=10):
text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
docs = text_splitter.split_documents(documents)
return docs
def initialize_model(documents):
with st.spinner("Processing documents may take a few seconds"):
new_pages = split_docs(documents)
if not new_pages:
st.error("No documents to process.")
return None
# db = Chroma.from_documents(new_pages, embedding_function)
db = Chroma.from_documents(
new_pages,
embedding_function,
persist_directory="../dataset"
)
db.persist()
# Use the new TogetherLLM class
llm = TogetherLLM()
retriever = db.as_retriever(similarity_score_threshold=0.95, search_kwargs={"k": 5})
prompt_template = """
CONTEXT: {context}
QUESTION: {question}"""
PROMPT = PromptTemplate(template=f"[INST] {prompt_template} [/INST]", input_variables=["context", "question"])
chain = RetrievalQA.from_chain_type(
llm=llm,
chain_type='stuff',
retriever=retriever,
input_key='query',
return_source_documents=True,
chain_type_kwargs={"prompt": PROMPT},
verbose=True
)
st.success("Document Processed successfully!")
return chain
class ConversationalAgent:
def __init__(self, chain):
self.chain = chain
self.history = []
def ask(self, query):
context = " ".join([item['response'] for item in self.history])
prompt_template = """
CONTEXT: {context}
QUESTION: {question}"""
prompt = f"[INST] CONTEXT: {context} QUESTION: {query} {prompt_template} [/INST]"
response = self.chain(query)
result = response['result']
# st.write(result)
self.history.append({'query': query, 'response': result})
return result, response['source_documents']
def process_file(uploaded_file):
"""Process the uploaded file and return an agent"""
documents = load_pdfs_from_file(uploaded_file)
if documents is None:
return None
chain = initialize_model(documents)
return ConversationalAgent(chain)
def demo_file_load():
dir_pre=os.getcwd()
pre=os.path.join(dir_pre,"dataset","LLM.pdf")
documents = load_pdfs_from_file(pre)
if documents is None:
return None
chain = initialize_model(documents)
return ConversationalAgent(chain) |