Update app.py

app.py

@@ -1,20 +1,15 @@
 import streamlit as st
-import re,torch
-import json,os
+import re, torch, json, os
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from datetime import datetime
-from huggingface_hub import login
+from huggingface_hub import login, InferenceClient
+import random
 
 login(token=os.getenv("TOKEN"))
 
-# Load model and tokenizer
-model = AutoModelForCausalLM.from_pretrained(
-    "google/gemma-2b",
-    torch_dtype="auto",
-    device_map="auto",
-)
+# Initialize the inference client for the Mixtral model
+client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
 if 'recipe' not in st.session_state:
     st.session_state.recipe = None
 
@@ -59,88 +54,18 @@ def create_detailed_prompt(user_direction, exclusions, serving_size, difficulty)
 
 def generate_recipe(user_inputs):
     with st.spinner('Building the perfect recipe...'):
-        # provide_recipe_schema = {
-        #     'type': 'function',
-        #     'function': {
-        #         'name': 'provide_recipe',
-        #         'description': 'Provides a detailed recipe strictly adhering to the user input/specifications, especially ingredient exclusions and the recipe difficulty',
-        #         'parameters': {
-        #             'type': 'object',
-        #             'properties': {
-        #                 'name': {
-        #                     'type': 'string',
-        #                     'description': 'A creative name for the recipe'
-        #                 },
-        #                 'description': {
-        #                     'type': 'string',
-        #                     'description': 'a brief one-sentence description of the provided recipe'
-        #                 },
-        #                 'ingredients': {
-        #                     'type': 'array',
-        #                     'items': {
-        #                         'type': 'object',
-        #                         'properties': {
-        #                             'name': {
-        #                                 'type': 'string',
-        #                                 'description': 'Quantity and name of the ingredient'
-        #                             }
-        #                         }
-        #                     }
-        #                 },
-        #                 'instructions': {
-        #                     'type': 'array',
-        #                     'items': {
-        #                         'type': 'object',
-        #                         'properties': {
-        #                             'step_number': {
-        #                                 'type': 'number',
-        #                                 'description': 'The sequence number of this step'
-        #                             },
-        #                             'instruction': {
-        #                                 'type': 'string',
-        #                                 'description': 'Detailed description of what to do in this step'
-        #                             }
-        #                         }
-        #                     }
-        #                 }
-        #             },
-        #             'required': [
-        #                 'name',
-        #                 'description',
-        #                 'ingredients',
-        #                 'instructions'
-        #             ]
-        #         }
-        #     }
-        # }
         prompt = create_detailed_prompt(user_inputs['user_direction'], user_inputs['exclusions'], user_inputs['serving_size'], user_inputs['difficulty'])
-        # messages = [{"role": "user", "content": prompt}]
-        # tool_section = "\n".join([f"{tool['function']['name']}({json.dumps(tool['function']['parameters'])})" for tool in [provide_recipe_schema]])
-        # text = f"{prompt}\n\nTools:\n{tool_section}"
         
-        # Tokenize and move to the correct device
-        model_inputs = tokenizer(prompt, return_tensors="pt")
-        # text = tokenizer.apply_chat_template(
-        #     messages,
-        #     tokenize=False,
-        #     add_generation_prompt=True,
-        #     tools=[provide_recipe_schema]
-        # )
-
-        # Tokenize and move to the correct device
-        # model_inputs = tokenizer([text], return_tensors="pt")
-        # torch.cuda.empty_cache()
-        # with torch.no_grad():
-        generated_ids = model.generate(
-            **model_inputs,
-            # max_new_tokens=512,
+        generate_kwargs = dict(
+            temperature=0.9,
+            max_new_tokens=1000,
+            top_p=0.9,
+            repetition_penalty=1.0,
+            do_sample=True,
         )
 
-        # generated_ids = [
-        #     output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-        # ]
-
-        st.session_state.recipe = tokenizer.decode(generated_ids[0])
+        response = client.text_generation(prompt, **generate_kwargs)
+        st.session_state.recipe = response['generated_text']
         st.session_state.recipe_saved = False
         
 def clear_inputs():
@@ -197,7 +122,7 @@ st.session_state.exclusions = st.text_area(
     placeholder="gluten, dairy, nuts, cilantro",
     )
 
-fancy_exclusions =""
+fancy_exclusions = ""
 
 if st.session_state.selected_difficulty == "Professional":
     exclude_fancy = st.checkbox(