Update app.py

fixed embeddings bug

app.py

@@ -14,15 +14,19 @@ writing_tones = ["Formal","Informal","Humorous","Serious","Sarcastic","Satirical
 
 
 # initialize client
+# we could try something larger, I need to check the models
 client = InferenceClient(
-    #"mistralai/Mixtral-8x7B-Instruct-v0.1" #wont fit onto free version
     "v1olet/v1olet_marcoroni-go-bruins-merge-7B"
 )
 
-# Load pre-trained mixtral tokenizer model (replace with your desired model if you want)
-model_id = "sentence-transformers/all-MiniLM-L6-v2" #using a small embeddings model"
+# Load pre-trained tokenizer model (replace with your desired model if you want, but it needs to be small)
+model_id = "sentence-transformers/all-MiniLM-L6-v2" #small embeddings model
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModel.from_pretrained(model_id) #load_in_4bit=True, use_flash_attention_2=True) #cant do this since HF free uses cpu
+model = AutoModel.from_pretrained(model_id) #cant do this since HF free uses cpu #load_in_4bit=True, 
+
+# Function to calculate cosine similarity between two embeddings
+def calculate_cosine_similarity(embedding1, embedding2):
+    return cosine_similarity(embedding1, embedding2)[0][0]
 
 # Function to convert text items into embeddings
 def get_embeddings(text_items):
@@ -35,40 +39,35 @@ def get_embeddings(text_items):
         embeddings.append(pooled_output)
     return embeddings
 
-# Helper Function to calculate cosine similarity between two embeddings
-# since we are using multiple variables, and since we want to compare more than just positive/negative ideas,
-# cosine similarity works better than euclidean distance as a measure of similarity
-def calculate_cosine_similarity(embedding1, embedding2):
-    return cosine_similarity(embedding1, embedding2)[0][0]
-
 # Helper Function to select values with small enough cosine similarity and concatenate them into a string
-def select_values_with_low_similarity(values, num_values_to_select, max_similarity):
+def select_values_with_low_similarity(embeddings, original_values, num_values_to_select, max_similarity):
     selected_values = []
     selected_indices = set()
 
     while len(selected_values) < num_values_to_select:
-        index1, index2 = random.sample(range(len(values)), 2)
-        item1, item2 = values[index1], values[index2]
+        index1, index2 = random.sample(range(len(embeddings)), 2)
+        embedding1, embedding2 = embeddings[index1], embeddings[index2]
 
-        if index1 != index2 and calculate_cosine_similarity(item1, item2) < max_similarity:
+        if index1 != index2 and calculate_cosine_similarity(embedding1, embedding2) < max_similarity:
             if index1 not in selected_indices:
-                selected_values.append(item1)
+                selected_values.append(original_values[index1])
                 selected_indices.add(index1)
             if index2 not in selected_indices:
-                selected_values.append(item2)
+                selected_values.append(original_values[index2])
                 selected_indices.add(index2)
 
     # Concatenate the selected values into a single string
     selected_string = ', '.join(selected_values)
     return selected_string
 
+
 # Convert text items into embeddings
 genre_embeddings = get_embeddings(book_genres)
 theme_embeddings = get_embeddings(book_themes)
 tone_embeddings = get_embeddings(writing_tones)
-#clear memory since we wont need to make the embeddings again
+#clear memory
 del model
-#torch.cuda.empty_cache() #only if using gpu
+#torch.cuda.empty_cache() 
 
 # helper function to format the prompt appropriately. 
 # For this creative writing tool, the user doesn't design the prompt itself 
@@ -76,11 +75,11 @@ del model
 def format_prompt(message, genres, tones, themes):
     # pick random ones if user leaves it blank but make sure they aren't opposites
     if not genres:
-        selected_genres = select_values_with_low_similarity(book_genres, random.randint(3, 5), 0.2)  # Adjust threshold as needed
+        selected_genres = select_values_with_low_similarity(genre_embeddings, book_genres, random.randint(3, 5), 0.2)  # Adjust threshold as needed
     if not tones:       
-        selected_tones = select_values_with_low_similarity(writing_tones, random.randint(3, 5), 0.2)  # Adjust threshold as needed
+        selected_tones = select_values_with_low_similarity(tone_embeddings, writing_tones, random.randint(3, 5), 0.2)  # Adjust threshold as needed
     if not themes:
-        selected_themes = select_values_with_low_similarity(book_themes, random.randint(3, 5), 0.2)  # Adjust threshold as needed
+        selected_themes = select_values_with_low_similarity(theme_embeddings, book_themes, random.randint(3, 5), 0.2)  # Adjust threshold as needed
 
     #Alpaca format since we can't use mixtral on free CPU settings
     prompt = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n"
@@ -93,7 +92,7 @@ def format_prompt(message, genres, tones, themes):
 
 
 # main function
-def generate(genres, themes, tones, system_prompt, temperature=1.25, max_new_tokens=256, top_p=0.95, repetition_penalty=1.15,):
+def generate(genres, themes, tones, system_prompt, temperature=1.25, max_new_tokens=512, top_p=0.95, repetition_penalty=1.15,):
     # check the temperature value, should not be too low, and make sure the values are floats
     temperature = float(temperature)
     if temperature < 1e-2:
@@ -111,7 +110,7 @@ def generate(genres, themes, tones, system_prompt, temperature=1.25, max_new_tok
     )
 
     formatted_prompt = format_prompt(f"{system_prompt}, '' ", genres, tones, themes)
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False, load_in_4bit=True, use_flash_attention_2=True)
     output = ""
 
     for response in stream:
@@ -166,21 +165,21 @@ additional_inputs=[
 
 def launch_interface():
     iface = gr.Interface(
-        fn=generate,
-        inputs=[
-            gr.Textbox("", label="Book Genres (comma-separated, or leave blank!)"),
-            gr.Textbox("", label="Book Themes (comma-separated, or leave blank!)"),
-            gr.Textbox("", label="Writing Tone (comma-separated, or leave blank!)"),
-            ],
-        additional_inputs=additional_inputs,
-        outputs="text",
-        live=False,
-        title="Novel Title and Summary Generator",
-        description='A fun creative writing tool, designed for when I have writer\'s block. Use it to practice building worlds, characters, scenes, etc. Write chapter 1, or a plot outline.' ,
-        theme='ParityError/Interstellar',
-    )
-
-    iface.launch(show_api=False)
+    fn=generate,
+    inputs=[
+      gr.Textbox("", label="Book Genres (comma-separated, or leave blank!)"),
+      gr.Textbox("", label="Book Themes (comma-separated, or leave blank!)"),
+      gr.Textbox("", label="Writing Tone (comma-separated, or leave blank!)"),
+      ],
+    #additional_inputs=additional_inputs,
+    outputs="text",
+    live=False,
+    title="Novel Title and Summary Generator",
+    description='A fun creative writing tool, designed for when I have writer\'s block. Use it to practice building worlds, characters, scenes, etc. Write chapter 1, or a plot outline.' ,
+    theme='ParityError/Interstellar', #)
+    additional_inputs=additional_inputs)
+    
+    iface.queue().launch(debug=True)
     
 if __name__=="__main__":
     launch_interface()