s

app.py

@@ -1,19 +1,14 @@
-# Requires gradio==3.16.2
-# Requires huggingface_hub==0.11.0
-
 import gradio as gr
-from huggingface_hub import InferenceClient
+import transformers
+import torch
 import logging
-from functools import lru_cache
 import html
 import signal
+from functools import lru_cache
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
 
-# Initialize the Hugging Face Inference Client
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
 # Constants
 MAX_HISTORY_LENGTH = 5  # Adjust as needed
 
@@ -46,34 +41,6 @@ def system_message_selector(choice, custom_message):
     else:
         return "You are a helpful assistant."
 
-@lru_cache(maxsize=32)
-def get_response_from_model(messages_tuple, max_tokens, temperature, top_p):
-    """
-    Calls the Hugging Face Inference API to get a response.
-
-    Parameters:
-        messages_tuple (tuple): A tuple of messages to be sent to the model.
-        max_tokens (int): Maximum number of tokens for the response.
-        temperature (float): Sampling temperature.
-        top_p (float): Top-p (nucleus) sampling parameter.
-
-    Returns:
-        str: The generated response from the model.
-    """
-    # Convert tuple back to list of dicts
-    messages = [dict(m) for m in messages_tuple]
-    response = ""
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-        response += token
-    return response
-
 def sanitize_input(text):
     """
     Sanitizes user input to prevent code injection or XSS attacks.
@@ -106,9 +73,55 @@ def validate_parameters(max_tokens, temperature, top_p):
         return False, "Error: 'Top-p' must be between 0.1 and 1.0."
     return True, ""
 
+# Load the model and tokenizer
+model_name = "HuggingFaceH4/zephyr-7b-beta"  # Replace with your actual model name
+
+try:
+    tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
+    model = transformers.AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float16,
+        device_map="auto"  # Automatically places model layers on available GPUs
+    )
+    model.eval()
+except Exception as e:
+    logging.error(f"Failed to load model {model_name}: {e}")
+    exit(1)
+
+@lru_cache(maxsize=32)
+def generate_response(prompt, max_tokens, temperature, top_p):
+    """
+    Generates a response using the loaded language model.
+
+    Parameters:
+        prompt (str): The input prompt for the model.
+        max_tokens (int): Maximum number of tokens for the response.
+        temperature (float): Sampling temperature.
+        top_p (float): Top-p (nucleus) sampling parameter.
+
+    Returns:
+        str: The generated response from the model.
+    """
+    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+    input_ids = input_ids.to(model.device)
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            input_ids,
+            max_length=input_ids.shape[1] + max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+
+    generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    return generated_text[len(prompt):].strip()
+
 def respond(message, history, persona_choice, custom_persona, max_tokens, temperature, top_p):
     """
-    Generates a response using the Hugging Face Inference API.
+    Generates a response using the loaded language model.
 
     Parameters:
         message (str): User's current input.
@@ -119,14 +132,13 @@ def respond(message, history, persona_choice, custom_persona, max_tokens, temper
         temperature (float): Sampling temperature.
         top_p (float): Top-p (nucleus sampling) parameter.
 
-    Yields:
+    Returns:
         str: The generated chatbot response.
     """
     # Validate parameters
     is_valid, error_message = validate_parameters(max_tokens, temperature, top_p)
     if not is_valid:
-        yield error_message
-        return
+        return error_message
 
     # Sanitize user input
     safe_message = sanitize_input(message)
@@ -138,33 +150,28 @@ def respond(message, history, persona_choice, custom_persona, max_tokens, temper
     # Select system message
     system_message = system_message_selector(persona_choice, custom_persona)
 
-    # Build messages with truncated history
-    messages = [{"role": "system", "content": system_message}]
+    # Build the conversation prompt
+    conversation = system_message + "\n\n"
     for user_msg, bot_msg in truncated_history:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if bot_msg:
-            messages.append({"role": "assistant", "content": bot_msg})
-    messages.append({"role": "user", "content": safe_message})
+        conversation += f"User: {user_msg}\n"
+        conversation += f"Assistant: {bot_msg}\n"
+    conversation += f"User: {safe_message}\nAssistant:"
 
     # Log the request
     logging.info(f"Received message: {safe_message}")
 
     try:
-        # Convert messages to a tuple of tuples for caching
-        messages_tuple = tuple(tuple(m.items()) for m in messages)
-
         # Use caching to optimize performance
-        response = get_response_from_model(
-            messages_tuple,
-            max_tokens,
-            temperature,
-            top_p,
+        response = generate_response(
+            prompt=conversation,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
         )
-        yield response
+        return response
     except Exception as e:
         logging.error(f"An error occurred: {e}")
-        yield "I'm sorry, but something went wrong. Please try again."
+        return "I'm sorry, but something went wrong. Please try again."
 
 # Create the UI components
 system_message_radio = gr.Radio(
@@ -179,7 +186,7 @@ system_message_textbox = gr.Textbox(
 )
 
 max_tokens_slider = gr.Slider(
-    minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"
+    minimum=1, maximum=1024, value=512, step=1, label="Max new tokens"
 )
 
 temperature_slider = gr.Slider(
@@ -190,7 +197,7 @@ top_p_slider = gr.Slider(
     minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"
 )
 
-# Create the ChatInterface directly with additional inputs
+# Create the ChatInterface
 demo = gr.ChatInterface(
     fn=respond,
     additional_inputs=[

requirements.txt

@@ -1,3 +1,3 @@
 transformers==4.31.0
 gradio==3.40.1
-torch==2.0.1
+torch==2.0.1