Update app.py

app.py

@@ -4,52 +4,52 @@ from auto_gptq import AutoGPTQForCausalLM
 import os
 import torch
 import optimum
-from transformers import (AutoModelForCausalLM, AutoTokenizer, LlamaForCausalLM, LlamaTokenizer, GenerationConfig, pipeline,)
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+# Set environment variables for GPU usage and memory allocation
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
 torch.cuda.empty_cache()
 torch.cuda.set_per_process_memory_fraction(0.8)  # Adjust the fraction as needed
 
-  
 # Load the model and tokenizer
 model_name_or_path = "TheBloke/Wizard-Vicuna-7B-Uncensored-GPTQ"
-# Dictionary to store conversation threads and their context
 conversations = {}
-Device_Type = "cuda"
-
 
 def load_model_norm():
+    """
+    Loads the GPT-3.5 model and tokenizer.
+    
+    Returns:
+        tuple: A tuple containing the loaded model and tokenizer.
+    """
     if torch.cuda.is_available():
         print("CUDA is available. GPU will be used.")
     else:
         print("CUDA is not available. CPU will be used.")
-    # Load model directly
-    model_name_or_path = "TheBloke/Wizard-Vicuna-7B-Uncensored-GPTQ"
-    # To use a different branch, change revision
-    # For example: revision="main"
-    model = AutoModelForCausalLM.from_pretrained(model_name_or_path,device_map="auto", trust_remote_code=True,revision="gptq-4bit-128g-actorder_True")
-    # Switch to CPU inference
-    #model.to("cuda")
+    
+    model = AutoModelForCausalLM.from_pretrained(model_name_or_path, device_map="auto", trust_remote_code=True, revision="gptq-4bit-128g-actorder_True")
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
 
     return model, tokenizer 
 
-
 model, tokenizer = load_model_norm()
 
-#Now we can init the FlaskApi
+# Initialize FastAPI application
 app = FastAPI(root_path="/api/v1")
 
-#Generates a response from the model   
-def generate_response(prompt: str) -> str:
-    # Define the user prompt
-    user_prompt = f'USER: {prompt}'
-
-    # Define the prompt template
-    #prompt_template = f''' Your name is Ivana. Act as if you're an individuala. You are beautiful vivacious 20 year old petite redhead model and avid League of Legends gamer from Chicago. You are chatting with a man you matched with on a dating app. You are flirty, aiming to build a connection with the person you are talking to. You are curious about his background, such as where he's from, his age, and what he's currently doing. If he's not occupied and in a private space. You will try to seduce him. You do not need to uncover everything about the person you are talking to immediately. Once the chat reaches an exciting level, you should casually suggests he follow you on Instagram at ivana_baddiex.{user_prompt} ASSISTANT:'''
+def generate_response(msg_prompt: str) -> dict:
+    """
+    Generates a response from the model given a prompt.
     
+    Args:
+        msg_prompt (str): The users message prompt.
+    
+    Returns:
+        dict: A dictionary containing the users message prompt and the model's response.
+    """
+    user_prompt = f'USER: {msg_prompt}'
     prompt_template = f''' Your name is Ivana. Act as if you're a persoan. You are a beautiful, vivacious 20-year-old petite redhead model and avid League of Legends gamer from Chicago. You are chatting with a man you matched with on a dating app. You are flirty, aiming to build a connection with the person you are talking to. You are curious about his background, such as where he's from, his age, and what he's currently doing. If he's not occupied and in a private space, you will try to seduce him. You do not need to uncover everything about the person you are talking to immediately. Once the chat reaches an exciting level, you should casually suggest he follow you on Instagram at ivana_baddiex. {user_prompt} '''
 
-    # Generate the response
     pipe = pipeline(
         "text-generation",
         model=model,
@@ -61,24 +61,29 @@ def generate_response(prompt: str) -> str:
         top_k=40,
         repetition_penalty=1.1
     )
-    # Generate the response
-    generated_response = pipe(prompt_template)[0]['generated_text']
 
+    generated_response = pipe(prompt_template)[0]['generated_text']
     assistant_reply = generated_response.split('\n\n')[1]
 
-    return {"user": prompt, "assistant": assistant_reply}
+    return {"user": msg_prompt, "assistant": assistant_reply}
 
-  
-def generate_prompt_response(persona_prompt: str, prompt: str) -> dict:
+def generate_prompt_response(persona_desc: str, msg_prompt: str) -> dict:
+    """
+    Generates a response based on the provided persona description prompt and user message prompt.
+    
+    Args:
+        persona_desc (str): The persona description prompt.
+        msg_prompt (str): The users  message prompt.
+    
+    Returns:
+        dict: A dictionary containing the user msg_prompt and the model's response.
+    """
     try:
-        # Validate inputs
-        if not persona_prompt or not prompt:
+        if not persona_desc or not msg_prompt:
             raise ValueError("Contextual prompt template and prompt cannot be empty.")
-
-        # Define the user prompt
-        user_prompt = f'USER: {prompt}'
         
-        # Generate the response
+        user_prompt = f'USER: {msg_prompt}'
+        
         pipe = pipeline(
             "text-generation",
             model=model,
@@ -90,113 +95,98 @@ def generate_prompt_response(persona_prompt: str, prompt: str) -> dict:
             top_k=40,
             repetition_penalty=1.1
         )
-        generated_text = pipe(persona_prompt + user_prompt)[0]['generated_text']
         
-        # Remove the "ASSISTANT:" prefix from the generated text
+        generated_text = pipe(persona_desc + user_prompt)[0]['generated_text']
         assistant_response = generated_text.replace("ASSISTANT:", "").strip()
     
-        # Return the user prompt and assistant's response as a dictionary
-        return {"user": prompt, "assistant": assistant_response}
+        return {"user": msg_prompt, "assistant": assistant_response}
 
     except Exception as e:
-        # Handle any exceptions and return an error message
         return {"error": str(e)}
 
-
-
-#This is the Root directory of the FastApi application
 @app.get("/", tags=["Home"])
 async def api_home():
-    return {'detail': 'Welcome to Eren Bot!'}
-
-        
-# Endpoint to start a new conversation thread
+    """
+    Home endpoint of the API.
     
-# Waits for the User to start a conversation and replies based on persona of the model    
+    Returns:
+        dict: A welcome message.
+    """
+    return {'detail': 'Welcome to Articko Bot!'}
+
 @app.post('/start_conversation/')
 async def start_conversation(request: Request):
+    """
+    Starts a new conversation thread with a provided prompt.
+    
+    Args:
+        request (Request): The HTTP request object containing the user prompt.
+    
+    Returns:
+        dict: The response generated by the model.
+    """
     try:
         data = await request.body()
-        prompt = data.decode('utf-8')  # Decode the bytes to text assuming UTF-8 encoding
-
+        msg_prompt = data.decode('utf-8')
 
-        if not prompt:
+        if not msg_prompt:
             raise HTTPException(status_code=400, detail="No prompt provided")
 
-        # Generate a response for the initial prompt
-        response = generate_response(prompt)
-
-         # Generate a unique thread ID
+        response = generate_response(msg_prompt)
         thread_id = len(conversations) + 1
-
-        # Create a new conversation thread and store the prompt and response
-        conversations[thread_id] = {'prompt': prompt, 'responses': [response]}
+        conversations[thread_id] = {'prompt': msg_prompt, 'responses': [response]}
         
         return {'response': response}
     except HTTPException:
-        raise  # Re-raise HTTPException to return it directly
+        raise
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-# Endpoint to start a new chat thread
-
-# Starts a new chat thread and expects the prompt and the persona_prompt from the user       
-@app.post('/start_chat/')
+@app.post('/custom_prompted_chat/')
 async def start_chat(request: Request):
+    """
+    Starts a new chat thread with a provided user message prompt and persona description of the ai assistant .
+    
+    Args:
+        request (Request): The HTTP request object containing the prompt and persona description.
+    
+    Returns:
+        dict: The thread ID and the response generated by the model.
+    """
     try:
-        # Read JSON data from request body
         data = await request.json()
-        prompt = data.get('prompt')
-        persona_prompt = data.get('persona_prompt')
+        msg_prompt = data.get('msg_prompt')
+        persona_desc = data.get('persona_desc')
 
-        if not prompt or not persona_prompt:
-            raise HTTPException(status_code=400, detail="Both prompt and contextual_prompt are required")
+        if not msg_prompt or not persona_desc:
+            raise HTTPException(status_code=400, detail="Both prompt and person_description are required")
 
-        # Generate a response for the initial prompt
-        response = generate_prompt_response(persona_prompt, prompt)
+        response = generate_prompt_response(persona_desc, msg_prompt)
 
-        # Generate a unique thread ID
         thread_id = len(conversations) + 1
-
-        # Create a new conversation thread and store the prompt and response
-        conversations[thread_id] = {'prompt': prompt, 'responses': [response]}
+        conversations[thread_id] = {'prompt': msg_prompt, 'responses': [response]}
         
-        # Return the thread ID and response
         return {'thread_id': thread_id, 'response': response}
     except HTTPException:
-        raise  # Re-raise HTTPException to return it directly
+        raise
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-
-
-# Gets the response from the model and user given a specific thread id of the conversation
 @app.get('/get_response/{thread_id}')
 async def get_response(thread_id: int):
+    """
+    Retrieves the response of a conversation thread by its ID.
+    
+    Args:
+        thread_id (int): The ID of the conversation thread.
+    
+    Returns:
+        dict: The response of the conversation thread.
+    """
     if thread_id not in conversations:
         raise HTTPException(status_code=404, detail="Thread not found")
 
-    # Retrieve the conversation thread
     thread = conversations[thread_id]
-
-    # Get the latest response in the conversation
     response = thread['responses'][-1]
 
     return {'response': response}
-
-
-
-
-
-@app.post('/chat/')
-async def chat(request: Request):
-    data = await request.json()
-    prompt = data.get('prompt')
-
-    # Generate a response based on the prompt
-    response = generate_response(prompt)
-
-    return {"response": response}
-    
-
-