Update app.py

app.py

@@ -1,14 +1,14 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import torch
+from huggingface_hub import InferenceClient
+import os
 
-# Load model and tokenizer
-model_name = "mistralai/Mistral-7B-Instruct-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
+# Initialize the Inference Client
+client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.1")
 
-# Create a text-generation pipeline
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=256)
+# Ensure you have set the HUGGINGFACE_TOKEN environment variable in your Hugging Face Space
+HF_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")
+if HF_TOKEN is None:
+    raise ValueError("Please set the HUGGINGFACE_TOKEN environment variable in your Hugging Face Space.")
 
 def format_prompt(message, history):
     fixed_prompt = """
@@ -83,18 +83,34 @@ def classify_mood(input_string):
             return word, True
     return None, False
 
-def generate(prompt, history):
+def generate(
+    prompt, history, temperature=0.7, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0,
+):
+    temperature = float(temperature)
+    if temperature < 1e-2:
+        temperature = 1e-2
+    top_p = float(top_p)
+
+    generate_kwargs = dict(
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        do_sample=True,
+    )
+
     formatted_prompt = format_prompt(prompt, history)
-    response = pipe(formatted_prompt, max_new_tokens=256, do_sample=True, temperature=0.7)[0]['generated_text']
-    
-    # Extract only the new generated text
-    new_text = response[len(formatted_prompt):].strip()
-    
-    mood, is_classified = classify_mood(new_text)
-    if is_classified:
-        playlist_message = f"Playing {mood.capitalize()} playlist for you!"
-        return playlist_message
-    return new_text
+
+    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+    output = ""
+
+    for response in stream:
+        output += response.token.text
+        mood, is_classified = classify_mood(output)
+        if is_classified:
+            playlist_message = f"Playing {mood.capitalize()} playlist for you!"
+            return playlist_message
+    return output
 
 def chat(message, history):
     response = generate(message, history)