Update app.py

app.py

@@ -1,133 +1,96 @@
-import shutil
 import os
-__import__('pysqlite3')
-import sys
-sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
-from sentence_transformers import SentenceTransformer
-import chromadb
-from datasets import load_dataset
 import gradio as gr
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
-
-# Set environment variables to address warnings
-os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
-
-
-torch.random.manual_seed(0)
-model_name = "microsoft/Phi-3-mini-4k-instruct-gguf"
+import copy
+from llama_cpp import Llama
+from huggingface_hub import hf_hub_download
+import chromadb
+from sentence_transformers import SentenceTransformer
 
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    low_cpu_mem_usage=True,
-    torch_dtype="auto",
-    trust_remote_code=True,
+# Initialize the Llama model
+llm = Llama(
+    model_path=hf_hub_download(
+        repo_id="microsoft/Phi-3-mini-4k-instruct-gguf",
+        filename="Phi-3-mini-4k-instruct-q4.gguf",
+    ),
+    n_ctx=2048,
+    n_gpu_layers=50,  # Adjust based on your VRAM
 )
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# Function to clear the cache
-def clear_cache(model_name):
-    cache_dir = os.path.expanduser(f'~/.cache/torch/sentence_transformers/{model_name.replace("/", "_")}')
-    if os.path.exists(cache_dir):
-        shutil.rmtree(cache_dir)
-        print(f"Cleared cache directory: {cache_dir}")
-    else:
-        print(f"No cache directory found for: {cache_dir}")
 
-# Embedding vector
+# Initialize ChromaDB Vector Store
 class VectorStore:
     def __init__(self, collection_name):
-        try:
-            self.embedding_model = SentenceTransformer('sentence-transformers/multi-qa-MiniLM-L6-cos-v1')
-        except Exception as e:
-            print(f"Error loading model: {e}")
-            raise
+        self.embedding_model = SentenceTransformer('sentence-transformers/multi-qa-MiniLM-L6-cos-v1')
         self.chroma_client = chromadb.Client()
         self.collection = self.chroma_client.create_collection(name=collection_name)
 
-    def populate_vectors(self, dataset, batch_size=20):
-        dataset = load_dataset('Thefoodprocessor/recipe_new_with_features_full', split='train')
-        dataset = dataset.select(range(1500))
-
-        texts = []
-        i = 0
-        for example in dataset:
-            title = example['title_cleaned']
-            recipe = example['recipe_new']
-            meal_type = example['meal_type']
-            allergy = example['allergy_type']
-            ingredients_alternative = example['ingredients_alternatives']
-            text = f"{title} {recipe} {meal_type} {allergy} {ingredients_alternative}"
-            texts.append(text)
-            if (i + 1) % batch_size == 0:
-                self._process_batch(texts, i)
-                texts = []
-            i += 1
-        if texts:
-            self._process_batch(texts, i)
-
-    def _process_batch(self, texts, batch_start_idx):
-        embeddings = self.embedding_model.encode(texts, batch_size=len(texts)).tolist()
-        for j, embedding in enumerate(embeddings):
-            self.collection.add(embeddings=[embedding], documents=[texts[j]], ids=[str(batch_start_idx + j)])
+    def populate_vectors(self, texts, ids):
+        embeddings = self.embedding_model.encode(texts, batch_size=32).tolist()
+        for text, embedding, doc_id in zip(texts, embeddings, ids):
+            self.collection.add(embeddings=[embedding], documents=[text], ids=[doc_id])
 
     def search_context(self, query, n_results=1):
-        query_embeddings = self.embedding_model.encode(query).tolist()
-        return self.collection.query(query_embeddings=query_embeddings, n_results=n_results)
+        query_embedding = self.embedding_model.encode([query]).tolist()
+        results = self.collection.query(query_embeddings=query_embedding, n_results=n_results)
+        return results['documents']
 
+# Example initialization (assuming you've already populated the vector store)
 vector_store = VectorStore("embedding_vector")
-vector_store.populate_vectors(dataset=None)
-
-def fine_tune_model():
-    dataset = load_dataset('Thefoodprocessor/recipe_new_with_features_full', split='train')
-    dataset = dataset.select(range(1500))
 
-    def tokenize_function(examples):
-        return tokenizer(
-            [" ".join([title, recipe]) for title, recipe in zip(examples['title_cleaned'], examples['recipe_new'])], 
-            padding="max_length", 
-            truncation=True
-        )
-
-    tokenized_datasets = dataset.map(tokenize_function, batched=True, batch_size=8)
-
-    training_args = TrainingArguments(
-        output_dir="./results",
-        evaluation_strategy="epoch",
-        learning_rate=2e-5,
-        per_device_train_batch_size=4,
-        per_device_eval_batch_size=4,
-        num_train_epochs=3,
-        weight_decay=0.01,
+# Populate with your data if not already done
+# vector_store.populate_vectors(your_texts, your_ids)
+
+def generate_text(
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+):
+    # Retrieve context from vector store
+    context_results = vector_store.search_context(message, n_results=1)
+    context = context_results[0] if context_results else ""
+
+    input_prompt = f"[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n {context}\n"
+    for interaction in history:
+        input_prompt += f"{interaction[0]} [/INST] {interaction[1]} </s><s> [INST] "
+    input_prompt += f"{message} [/INST] "
+
+    temp = ""
+    output = llm(
+        input_prompt,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=40,
+        repeat_penalty=1.1,
+        max_tokens=max_tokens,
+        stop=["", " \n", "ASSISTANT:", "USER:", "SYSTEM:"],
+        stream=True,
     )
+    for out in output:
+        temp += out["choices"][0]["text"]
+        yield temp
+
+# Define the Gradio interface
+demo = gr.ChatInterface(
+    generate_text,
+    title="llama-cpp-python on GPU with ChromaDB",
+    description="Running LLM with context retrieval from ChromaDB",
+    examples=[
+        ["I have leftover rice, what can I make out of it?"],
+        ["Can I make lunch for two people with this?"],
+    ],
+    cache_examples=False,
+    retry_btn=None,
+    undo_btn="Delete Previous",
+    clear_btn="Clear",
+    additional_inputs=[
+        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
+    ],
+)
 
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=tokenized_datasets,
-    )
-
-    trainer.train()
-
-fine_tune_model()
-
-conversation_history = []
-
-def chatbot_response(user_input):
-    global conversation_history
-    results = vector_store.search_context(user_input, n_results=1)
-    context = results['documents'][0] if results['documents'] else ""
-    conversation_history.append(f"User: {user_input}\nContext: {context[:150]}\nBot:")
-    inputs = tokenizer("\n".join(conversation_history), return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=150, do_sample=True, temperature=0.7)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    conversation_history.append(response)
-    return response
-
-def chat(user_input):
-    response = chatbot_response(user_input)
-    return response
-
-css = ".gradio-container {background: url(https://upload.wikimedia.org/wikipedia/commons/f/f5/Spring_Kitchen_Line-Up_%28Unsplash%29.jpg)}"
-iface = gr.Interface(fn=chat, inputs="text", outputs="text", css=css)
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()