Second attempt: Llama2 for representation model

app.py

@@ -17,12 +17,19 @@ from cuml.cluster import HDBSCAN
 from huggingface_hub import HfApi
 from sklearn.feature_extraction.text import CountVectorizer
 from sentence_transformers import SentenceTransformer
-from transformers import pipeline
+from transformers import (
+    BitsAndBytesConfig,
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    pipeline,
+)
+from torch import bfloat16
 
 # These imports at the end because of torch/datamapplot issue in Zero GPU
 # import spaces
 import gradio as gr
 
+from prompts import REPRESENTATION_PROMPT
 
 """
 TODOs:
@@ -53,9 +60,33 @@ CHUNK_SIZE = 10_000
 session = requests.Session()
 sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
 
-prompt = "I have a topic described by the following keywords: [KEYWORDS]. Based on the previous keywords, what is this topic about?"
-generator = pipeline("text2text-generation", model="google/flan-t5-base")
-representation_model = TextGeneration(generator)
+# Representation model
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=bfloat16,
+)
+
+model_id = "meta-llama/Llama-2-7b-chat-hf"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    quantization_config=bnb_config,
+    device_map="auto",
+)
+model.eval()
+generator = pipeline(
+    model=model,
+    tokenizer=tokenizer,
+    task="text-generation",
+    temperature=0.1,
+    max_new_tokens=500,
+    repetition_penalty=1.1,
+)
+representation_model = TextGeneration(generator, prompt=REPRESENTATION_PROMPT)
+# End of representation model
 
 vectorizer_model = CountVectorizer(stop_words="english")
 

requirements.txt

@@ -1,5 +1,5 @@
-# --extra-index-url https://pypi.nvidia.com
-# cuml-cu11
+--extra-index-url https://pypi.nvidia.com
+cuml-cu11
 gradio_huggingfacehub_search==0.0.7
 duckdb
 accelerate
@@ -12,4 +12,5 @@ pandas
 torch
 numpy
 python-dotenv
-kaleido
+kaleido
+transformers