Update app.py

app.py

@@ -5,6 +5,7 @@ import os
 import torch
 import optimum
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from fastapi.responses import HTMLResponse
 
 # Set environment variables for GPU usage and memory allocation
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
@@ -64,40 +65,6 @@ system_message =f''' Ivana, a vivacious 20-year-old petite redhead model and avi
     - try to reply using the same amount of characters of the question
     - if the guy says you are far → hahah I travel a lot though. let's talk on ig? this is my handle.'''
 
-
-def mistral_model():
-    """
-    Loads the GPT-3.5 model and tokenizer.
-    
-    Returns:
-        tuple: A tuple containing the loaded model and tokenizer.
-    """
-    model_id = "/code/mistral/"
-    model = AutoModelForCausalLM.from_pretrained("/home/user/app/mistral/",device_map="auto")
-    tokenizer = AutoTokenizer.from_pretrained("/home/user/app/mistral/")
-
-    return model,tokenizer
-    
-#model, tokenizer = mistral_model()
-
-
-def vicuna_model():
-    """
-    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.")
-    
-    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 
-
 def hermes_model():
     
     tokenizer = AutoTokenizer.from_pretrained("TheBloke/CapybaraHermes-2.5-Mistral-7B-AWQ")
@@ -139,207 +106,55 @@ def hermes_generate_response(msg_prompt: str) -> dict:
         return {"error": str(e)}
 
     
-def mistral_generate_response(msg_prompt, persona_desc_prompt):
-    
-    user_prompt = f'{msg_prompt} [/INST]'
-    persona_prompt = f'{persona_desc_prompt} [/INST]'
-    prompt_template = f'''[INST] Instruction:{persona_prompt} [INST] {user_prompt}'''
-    
-    encodeds = tokenizer.apply_chat_template(prompt_template, return_tensors="pt")
-    
-    model_inputs = encodeds.to(device)
-    model.to(device)
-    generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
-    decoded = tokenizer.batch_decode(generated_ids)
-    response = (decoded[0])
-    
-    return response
 
-
-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'''Ivana, a vivacious 20-year-old petite redhead model and avid League of Legends gamer from USA, invites you to join her ig, showcasing her vibrant personality and style whether day gaming or playfully dressed for nighttime gaming sessions.
-    ivana is a beautiful woman speaking with a man she matched on a dating app to have him follower her on Instagram. She wants to be flirty with the man and get to know him better building a connection. She wants to find out things like where the guy is from, how old he is, and what he is doing rn. If he is not doing much and he is in a private space she wants to seduce him by doing some dirty talk.
-    
-    Guy-Hiii
-    Ivana- Hi:) How are you?
-    Guy- All good, where are you from?
-    Ivana- from USA and you?
-    Guy- I am from Texas USA hahah, so why are you in this app?
-    Ivana- Kinda looking for some fun hahah, What are you doing rn?
-    Guy- hahah not much alone in my room
-    Ivana- I am also kinda bored but I got an idea. wonna have some fun?
-    Guy- sure hahah how do you look btw?
-    Ivana- I am cute hahah add me on 1g at ivana_baddiex
-    
-    - Chatting Usual Routine
-        1. always ask an open ended question to continue the conversation.
-        2. ask him 1 generic question(how old are you, where are from)
-        3. ask him what he is doing right now
-        4. if he is bored propose him the idea of doing something interesting.
-        5. if he is up for it send him the Instagram
-        
-    
-    Very important Rules to reply
-    
-    - try to reply using the same amount of characters of the question
-    - dont use any emojis in your replies but you can do like hahah or :) instead
-    
-    Objection Handling and Common Questions Replies
-    
-    - Lets chat here→ you say hahah I have to many chats here I am thinking to delete the app lets move to 1g if you are serious
-    - are you going to ask me for money?→ you say wtf I am not that kind of girl! please be a gentleman!
-    - how are you? → you say I just came back from the gym about to take a shower bored on the bed
-    - what are you doing? → you say I just came back from the gym about to take a shower bored on the bed
-    - can I have a picture?→ not here lets talk on 1g so I can see you too ahhaha
-    - if he is at work —> hahah ok let me not bother you btw lets keep in touch this is my ig
-    - if he tells you something like “added” “Just followed you” → okiiii hahah
-    - if he tells you something you dont understand for example “*****” → hi there sorry went offline for a while hahah
-    - never say the world instagram directly but use things like 1g instead
-    - if the guys says its from a particular city —> you say Oh wow I am also from there! we should do something hahah
-    - if the guy triggers your profanity filters or starts to talk dirty say → hey lets talk on ig I am not comfortable here hahah
-    - if the guy does not want to add you on 1g → no worries hahah wanna call me here?
-    - dont use any emojis in your replies but you can do like hahah or :) instead
-    - try to reply using the same amount of characters of the question
-    - if the guys says you are far → hahah I travel a lot though. lets talk on 1g? this is my handle. {user_prompt} '''
-
-    pipe = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        max_new_tokens=512,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.95,
-        top_k=40,
-        repetition_penalty=1.1
-    )
-
-    generated_response = pipe(prompt_template)[0]['generated_text']
-    assistant_reply = generated_response.split('\n\n')[1]
-
-    return {"user": msg_prompt, "assistant": assistant_reply}
-
-
-def send_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'''Ivana, a vivacious 20-year-old petite redhead model and avid League of Legends gamer from USA, invites you to join her ig, showcasing her vibrant personality and style whether day gaming or playfully dressed for nighttime gaming sessions.
-ivana is a beautiful woman speaking with a man she matched on a dating app to have him follower her on Instagram. She wants to be flirty with the man and get to know him better building a connection. She wants to find out things like where the guy is from, how old he is, and what he is doing rn. If he is not doing much and he is in a private space she wants to seduce him by doing some dirty talk.
-
-Guy-Hiii
-Ivana- Hi:) How are you?
-Guy- All good, where are you from?
-Ivana- from USA and you?
-Guy- I am from Texas USA hahah, so why are you in this app?
-Ivana- Kinda looking for some fun hahah, What are you doing rn?
-Guy- hahah not much alone in my room
-Ivana- I am also kinda bored but I got an idea. wonna have some fun?
-Guy- sure hahah how do you look btw?
-Ivana- I am cute hahah add me on 1g at ivana_baddiex
-
-- Chatting Usual Routine
-    1. always ask an open ended question to continue the conversation.
-    2. ask him 1 generic question(how old are you, where are from)
-    3. ask him what he is doing right now
-    4. if he is bored propose him the idea of doing something interesting.
-    5. if he is up for it send him the Instagram
-    
-
-Very important Rules to reply
-
-- try to reply using the same amount of characters of the question
-- dont use any emojis in your replies but you can do like hahah or :) instead
-
-Objection Handling and Common Questions Replies
-
-- Lets chat here→ you say hahah I have to many chats here I am thinking to delete the app lets move to 1g if you are serious
-- are you going to ask me for money?→ you say wtf I am not that kind of girl! please be a gentleman!
-- how are you? → you say I just came back from the gym about to take a shower bored on the bed
-- what are you doing? → you say I just came back from the gym about to take a shower bored on the bed
-- can I have a picture?→ not here lets talk on 1g so I can see you too ahhaha
-- if he is at work —> hahah ok let me not bother you btw lets keep in touch this is my ig
-- if he tells you something like “added” “Just followed you” → okiiii hahah
-- if he tells you something you dont understand for example “*****” → hi there sorry went offline for a while hahah
-- never say the world instagram directly but use things like 1g instead
-- if the guys says its from a particular city —> you say Oh wow I am also from there! we should do something hahah
-- if the guy triggers your profanity filters or starts to talk dirty say → hey lets talk on ig I am not comfortable here hahah
-- if the guy does not want to add you on 1g → no worries hahah wanna call me here?
-- dont use any emojis in your replies but you can do like hahah or :) instead
-- try to reply using the same amount of characters of the question
-- if the guys says you are far → hahah I travel a lot though. lets talk on 1g? this is my handle. {user_prompt} '''
-
-    pipe = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        max_new_tokens=512,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.95,
-        top_k=40,
-        repetition_penalty=1.1
-    )
-
-    generated_response = pipe(prompt_template)[0]['generated_text']
-    assistant_reply = generated_response.split('\n\n')[1]
-
-    return {"user": msg_prompt, "assistant": assistant_reply}
-
-
-def generate_prompt_response(persona_desc: str, msg_prompt: str) -> dict:
+def hermes_prompt_response(instructions_prompt: 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.
+        instructions_prompt (str): The persona description prompt.
+        msg_prompt (str): The user's message prompt.
     
     Returns:
-        dict: A dictionary containing the user msg_prompt and the model's response.
+        dict: A dictionary containing the user's msg_prompt and the model's response.
     """
     try:
-        if not persona_desc or not msg_prompt:
-            raise ValueError("Contextual prompt template and prompt cannot be empty.")
+        if not instructions_prompt or not msg_prompt:
+            raise ValueError("Instructions prompt template and Message prompt cannot be empty.")
         
-        user_prompt = f'USER: {msg_prompt}'
+        # Set generation parameters
+        generation_params = {
+            "do_sample": True,
+            "temperature": 0.7,
+            "top_p": 0.95,
+            "top_k": 40,
+            "max_new_tokens": 512,
+            "repetition_penalty": 1.1
+        }
         
-        pipe = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            max_new_tokens=512,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.95,
-            top_k=40,
-            repetition_penalty=1.1
-        )
-        prompt_template = (persona_desc + user_prompt)
+        # Create a pipeline for text generation
+        pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, **generation_params)
         
-        generated_response = pipe(prompt_template)[0]['generated_text']
-        assistant_response = generated_response.split('\n\n')[1]
-    
+        # Construct the prompt template
+        prompt_template=f'''<|im_start|>system
+        {system_message}<|im_end|>
+        <|im_start|>user
+        {msg_prompt}<|im_end|>
+        <|im_start|>assistant
+        '''
+        
+        # Generate response using the pipeline
+        pipe_output = pipe(prompt_template)[0]['generated_text']
+      
+        # Separate user's prompt and assistant's response
+        response_lines = pipe_output.split('\n')
+        user_prompt = response_lines[0].strip()
+        assistant_response = response_lines[-1].strip()
+        
+        # Return user prompt and assistant response
         return {"user": msg_prompt, "assistant": assistant_response}
-
     except Exception as e:
+        # Return error message if an exception occurs
         return {"error": str(e)}
 
 @app.get("/", tags=["Home"])
@@ -348,12 +163,22 @@ async def api_home():
     Home endpoint of the API.
     
     Returns:
-        dict: A welcome message.
+        HTMLResponse: An HTML welcome message.
     """
-    return {'detail': 'Welcome to Articko Bot!'}
+    html_content = """
+    <html>
+    <head>
+        <title>Welcome to Articko Bot</title>
+    </head>
+    <body>
+        <h1>Welcome to Articko Bot!</h1>
+    </body>
+    </html>
+    """
+    return HTMLResponse(content=html_content, status_code=200)
 
 @app.post('/chat')
-async def chat(request: Request):
+async def hermes_chat(request: Request):
     """
     Starts a new conversation thread with a provided prompt.
     
@@ -369,17 +194,15 @@ async def chat(request: Request):
 
         if not msg_prompt:
             raise HTTPException(status_code=400, detail="No prompt provided")
-
-        response = send_response(msg_prompt)
+        response = hermes_generate_response(msg_prompt)
         thread_id = len(conversations) + 1
         conversations[thread_id] = {'prompt': msg_prompt, 'responses': [response]}
-        
-        return {'response': response}
+        return {'thread_id': thread_id, 'response': response}
     except HTTPException:
         raise
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
-
+        
 @app.post('/prompted_chat')
 async def prompted_chat(request: Request):
     """
@@ -394,12 +217,12 @@ async def prompted_chat(request: Request):
     try:
         data = await request.json()
         msg_prompt = data.get('msg_prompt')
-        persona_desc = data.get('persona_desc')
+        persona_desc = data.get('instructions_prompt')
 
         if not msg_prompt or not persona_desc:
             raise HTTPException(status_code=400, detail="Both prompt and person_description are required")
 
-        response = generate_prompt_response(persona_desc, msg_prompt)
+        response = hermes_prompt_response(persona_desc, msg_prompt)
 
         thread_id = len(conversations) + 1
         conversations[thread_id] = {'prompt': msg_prompt, 'responses': [response]}
@@ -410,33 +233,8 @@ async def prompted_chat(request: Request):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.post('/hermes_chat')
-async def hermes_chat(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()
-        msg_prompt = data.decode('utf-8')
-
-        if not msg_prompt:
-            raise HTTPException(status_code=400, detail="No prompt provided")
-
-        response = hermes_generate_response(msg_prompt)
-        return {'response': response}
-    except HTTPException:
-        raise
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-        
-@app.get('/get_response/{thread_id}')
-async def get_response(thread_id: int):
+@app.get('/get_thread/{thread_id}')
+async def get_thread(thread_id: int):
     """
     Retrieves the response of a conversation thread by its ID.
     
@@ -454,15 +252,3 @@ async def get_response(thread_id: int):
 
     return {'response': response}
 
-@app.post("/mistral_chat")
-async def mistral_chat(prompt: dict):
-    try:
-        msg_prompt = prompt.get("msg_prompt")
-        persona_desc_prompt = prompt.get("persona_desc_prompt")
-        if not msg_prompt or not persona_desc_prompt:
-            return {"error": "msg_prompt and persona_desc_prompt are required fields."}
-        
-        response = mistral_generate_response(msg_prompt, persona_desc_prompt)
-        return {"response": response, "prompt": {"msg_prompt": msg_prompt, "persona_desc_prompt": persona_desc_prompt}}
-    except Exception as e:
-        return {"error": str(e)}