Update app.py

app.py

@@ -8,21 +8,25 @@ 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'
 
+# Ensure necessary packages are installed
+!pip install accelerate
+!pip install flash-attention
 
-# Load model directly
-from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
-os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
-torch.random.manual_seed(0) 
-model = AutoModelForCausalLM.from_pretrained( 
-    "microsoft/Phi-3-mini-4k-instruct",  
-    low_cpu_mem_usage=True,  
-    torch_dtype="auto",  
-    trust_remote_code=True,  
-) 
+torch.random.manual_seed(0)
+model_name = "microsoft/Phi-3-mini-4k-instruct"
 
-tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct") 
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    low_cpu_mem_usage=True,
+    torch_dtype="auto",
+    trust_remote_code=True,
+)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Function to clear the cache
 def clear_cache(model_name):
@@ -36,7 +40,6 @@ def clear_cache(model_name):
 # Embedding vector
 class VectorStore:
     def __init__(self, collection_name):
-        # Initialize the embedding model
         try:
             self.embedding_model = SentenceTransformer('sentence-transformers/multi-qa-MiniLM-L6-cos-v1')
         except Exception as e:
@@ -45,33 +48,24 @@ class VectorStore:
         self.chroma_client = chromadb.Client()
         self.collection = self.chroma_client.create_collection(name=collection_name)
 
-    # Method to populate the vector store with embeddings from a dataset
-    def populate_vectors(self, dataset, batch_size=10):
-        # Use dataset streaming
-        #dataset = load_dataset('Thefoodprocessor/recipe_new_with_features_full', split='train[:1500]', streaming=True)
+    def populate_vectors(self, dataset, batch_size=20):
         dataset = load_dataset('Thefoodprocessor/recipe_new_with_features_full', split='train')
-        dataset = dataset.select(range(50))  # Select the first 1500 examples
+        dataset = dataset.select(range(1500))
 
         texts = []
-        i = 0  # Initialize index
+        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']
-
-            # Concatenate the text from the columns
-            text = f"{title} {recipe} {allergy}"
+            ingredients_alternative = example['ingredients_alternatives']
+            text = f"{title} {recipe} {meal_type} {allergy} {ingredients_alternative}"
             texts.append(text)
-
-            # Process the batch
             if (i + 1) % batch_size == 0:
                 self._process_batch(texts, i)
                 texts = []
-
-            i += 1  # Increment index
-
-        # Process the remaining texts
+            i += 1
         if texts:
             self._process_batch(texts, i)
 
@@ -84,47 +78,42 @@ class VectorStore:
         query_embeddings = self.embedding_model.encode(query).tolist()
         return self.collection.query(query_embeddings=query_embeddings, n_results=n_results)
 
-# Create a vector embedding
 vector_store = VectorStore("embedding_vector")
 vector_store.populate_vectors(dataset=None)
 
-# Fine-tuning function
 def fine_tune_model():
-    # Load your dataset
     dataset = load_dataset('Thefoodprocessor/recipe_new_with_features_full', split='train')
-    dataset = dataset.select(range(50))  # Select the first 1500 examples
+    dataset = dataset.select(range(1500))
 
-    # Prepare the data for training
     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)
+        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)
+    tokenized_datasets = dataset.map(tokenize_function, batched=True, batch_size=8)
 
-    # Define training arguments
     training_args = TrainingArguments(
         output_dir="./results",
         evaluation_strategy="epoch",
         learning_rate=2e-5,
-        per_device_train_batch_size=8,
-        per_device_eval_batch_size=8,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
         num_train_epochs=3,
         weight_decay=0.01,
     )
 
-    # Initialize Trainer
     trainer = Trainer(
         model=model,
         args=training_args,
         train_dataset=tokenized_datasets,
     )
 
-    # Train the model
     trainer.train()
 
-# Fine-tune the model
 fine_tune_model()
 
-# Define the chatbot response function
 conversation_history = []
 
 def chatbot_response(user_input):
@@ -138,7 +127,6 @@ def chatbot_response(user_input):
     conversation_history.append(response)
     return response
 
-# Gradio interface
 def chat(user_input):
     response = chatbot_response(user_input)
     return response