Create app.py

app.py

@@ -0,0 +1,103 @@
+from langchain.embeddings.huggingface import HuggingFaceEmbeddings
+import os, pincone, time, transformers
+from datasets import load_dataset
+from torch import bfloat16
+from langchain.llms import HuggingFacePipeline
+from langchain.vectorstores import Pinecone
+from langchain.chains import RetrievalQA
+
+embed_model_id = 'sentence-transformers/all-MiniLM-L6-v2'
+device = 'cpu'
+
+embed_model = HuggingFaceEmbeddings(
+    model_name=embed_model_id,
+    model_kwargs={'device': device},
+    encode_kwargs={'device': device, 'batch_size': 32}
+)
+
+api_key = os.environ.get('PINCONE_API_KEY')
+env_name = os.environ.get('PINECONE_ENV')
+
+pinecone.init(
+    api_key=api_key, 
+    environment=env_name
+)
+
+index_name = 'llama-2-rag-2.0'
+
+if index_name not in pinecone.list_indexes():
+    pinecone.create_index(
+        index_name,
+        dimension=len(embeddings[0]),
+        metric='cosine'
+    )
+    while not pinecone.describe_index(index_name).status['ready']:
+        time.sleep(1)
+
+index = pinecone.Index(index_name)
+
+data = load_dataset('jamescalam/llama-2-arxiv-papers-chunked', split='train')
+data = data.to_pandas()
+batch_size = 32
+
+for i in range(0, len(data), batch_size):
+    i_end = min(len(data), i+batch_size)
+    batch = data.iloc[i:i_end]
+    ids = [f"{x['doi']}-{x['chunk-id']}" for i, x in batch.iterrows()]
+    texts = [x['chunk'] for i, x in batch.iterrows()]
+    embeds = embed_model.embed_documents(texts)
+    metadata = [
+        {'text': x['chunk'],
+        'source': x['source'],
+        'title': x['title']} for i, x in batch.iterrows()
+    ]
+    index.upsert(vectors=zip(ids, embeds, metadata))
+
+model_id = 'meta-llama/Llama-2-7b-chat-hf'
+
+bnb_config = transformers.BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type='nf4',
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=bfloat16
+)
+model_config = transformers.AutoConfig.from_pretrained(model_id)
+model = transformers.AutoModelForCausalLM.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    config=model_config,
+    quantization_config=bnb_config,
+    device_map='auto'
+)
+model.eval()
+
+tokenizer = transformer.AutoTokenizer.from_pretrained(model_id)
+
+generate_text = transformers.pipeline(
+    model=model,
+    tokenizer=tokenizer,
+    return_full_text=True,
+    task='text-generation',
+    temperature=0.3,
+    max_new_tokens=512,
+    repetition_penalty=1.1
+)
+llm = HuggingFacePipeline(pipeline=generate_text)
+text_field = 'text'
+vectorstore = Pinecone(index, embed_model.embed_query, text_field)
+rag_pipeline = RetrievalQA.from_chain_type(
+    llm=llm,
+    chain_type='stuff',
+    retriever=vectorstore.as_retriever()
+)
+
+title = 'arxiv-retrieval'
+
+def predict(input):
+    return rag_pipeline(input)
+
+gr.Interface(
+    fn=predict,
+    inputs=['text', 'state'],
+    outputs=['chatbot', 'state']
+).launch()