Update app.py

app.py

@@ -7,75 +7,87 @@ def load_llm():
     Loads the GPT-2 model and tokenizer using the Hugging Face `transformers` library.
     """
     try:
-        print("Loading GPT-2 model and tokenizer...")
-        model_name = 'gpt2'  # Replace with your custom model name if using a fine-tuned version
+        print("Downloading or loading the GPT-2 model and tokenizer...")
+        model_name = 'gpt2'  # Replace with your custom model if available
         model = GPT2LMHeadModel.from_pretrained(model_name)
         tokenizer = GPT2Tokenizer.from_pretrained(model_name)
         print("Model and tokenizer successfully loaded!")
         return model, tokenizer
     except Exception as e:
-        print(f"Error during model loading: {e}")
+        print(f"An error occurred while loading the model: {e}")
         return None, None
 
-def generate_response(model, tokenizer, user_input, max_length=512):
+def generate_response(model, tokenizer, user_input):
     """
-    Generates a response using the GPT-2 model based on user input.
+    Generates a response using the GPT-2 model and tokenizer.
     
     Args:
-    - model: The GPT-2 model.
-    - tokenizer: The corresponding tokenizer.
-    - user_input (str): The user's input message.
-    - max_length (int): The maximum length of the generated output.
+    - model: The loaded GPT-2 model.
+    - tokenizer: The tokenizer corresponding to the GPT-2 model.
+    - user_input (str): The input question from the user.
 
     Returns:
-    - response (str): The AI-generated response.
+    - response (str): The generated response.
     """
     try:
         inputs = tokenizer.encode(user_input, return_tensors='pt')
-        outputs = model.generate(inputs, max_length=max_length, num_return_sequences=1)
+        outputs = model.generate(inputs, max_length=512, num_return_sequences=1)
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return response
     except Exception as e:
-        return f"Error during response generation: {e}"
+        return f"An error occurred during response generation: {e}"
 
 # Load the model and tokenizer
 model, tokenizer = load_llm()
 
 if model is None or tokenizer is None:
-    print("Failed to load model and tokenizer.")
+    print("Model and/or tokenizer loading failed.")
 else:
-    print("Model and tokenizer are ready to use.")
+    print("Model and tokenizer are ready for use.")
 
-# Initialize the Hugging Face API client
+# Initialize the Hugging Face API client (ensure it’s correctly set up)
 client = InferenceClient()
 
 def respond(message, history, system_message, max_tokens, temperature, top_p):
     """
-    Handles the chatbot interaction, sending conversation history and system message
-    to the Hugging Face Inference API for generating AI responses.
+    Handles interaction with the chatbot by sending the conversation history
+    and system message to the Hugging Face Inference API.
     """
-    print("Respond function initiated")
-    print("User message:", message)
-    print("Chat history:", history)
+    print("Starting respond function")
+    print("Received message:", message)
+    print("Conversation history:", history)
 
-    # Construct the conversation history with the system message
     messages = [{"role": "system", "content": system_message}]
     
     for user_msg, assistant_msg in history:
         if user_msg:
-            print("Adding user message to history:", user_msg)
+            print("Adding user message to messages:", user_msg)
             messages.append({"role": "user", "content": user_msg})
         if assistant_msg:
-            print("Adding assistant message to history:", assistant_msg)
+            print("Adding assistant message to messages:", assistant_msg)
             messages.append({"role": "assistant", "content": assistant_msg})
     
     messages.append({"role": "user", "content": message})
-    print("Message list for model:", messages)
+    print("Final message list for the model:", messages)
 
-    # Use GPT-2 model for local generation
-    conversation_history = " ".join([f"{msg['role']}: {msg['content']}" for msg in messages])
-    response = generate_response(model, tokenizer, conversation_history, max_length=max_tokens)
-    return response
+    response = ""
+    try:
+        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
+            print("Token received:", token)
+            yield response
+    except Exception as e:
+        print("An error occurred:", e)
+        yield f"An error occurred: {e}"
+
+    print("Response generation completed")
 
 # Set up the Gradio ChatInterface
 demo = gr.ChatInterface(