app.py

import os
import torch
from flask import Flask, jsonify, request
from flask_cors import CORS
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, GenerationConfig
from accelerate import Accelerator
import re
import traceback
from transformers import pipeline 
from sentence_transformers import SentenceTransformer, util
  
# Set the HF_HOME environment variable to a writable directory
os.environ["HF_HOME"] = "/workspace/huggingface_cache"

app = Flask(__name__)

# Enable CORS for specific origins
CORS(app, resources={r"/send_message": {"origins": ["http://localhost:3000", "https://main.dbn2ikif9ou3g.amplifyapp.com"]}})
  
# Load zero-shot classification pipeline
classifier = pipeline("zero-shot-classification")

 # Load Sentence-BERT model
bertmodel = SentenceTransformer('all-MiniLM-L6-v2')  # Lightweight, efficient model; choose larger if needed
        
# Global variables for model and tokenizer
model = None
tokenizer = None
accelerator = Accelerator() 
highest_label = None 
loaded_models = {}
 
def get_model_and_tokenizer(model_id: str):
    """
    Load and cache the model and tokenizer for the given model_id.
    """
    global model, tokenizer  # Declare global variables to modify them within the function
    if model_id not in loaded_models:
        try:
            tokenizer = AutoTokenizer.from_pretrained(model_id)
            model = AutoModelForCausalLM.from_pretrained(model_id)
            model = accelerator.prepare(model)
            loaded_models[model_id] = (model, tokenizer)
        except Exception as e:
            print("Error loading model:")
            print(traceback.format_exc())  # Logs the full error traceback
            raise e  # Reraise the exception to stop execution
    return loaded_models[model_id]
        
        
# Extract the core sentence needing grammar correction
def extract_core_sentence(user_input):
    """
    Extract the core sentence needing grammar correction from the user input.
    """
    match = re.search(r"(?<=sentence[: ]).+", user_input, re.IGNORECASE)
    if match:
        return match.group(0).strip()
    return user_input

def classify_intent(user_input):
    """
    Classify the intent of the user input using zero-shot classification.
    """
    candidate_labels = [
        "grammar correction", "information request", "task completion", 
        "dialog continuation", "personal opinion", "product inquiry",
        "feedback request", "recommendation request", "clarification request", 
        "affirmation or agreement", "real-time data request", "current information"
    ]
    result = classifier(user_input, candidate_labels)
    highest_score_index = result['scores'].index(max(result['scores']))
    highest_label = result['labels'][highest_score_index]
    return highest_label


# Reformulate the prompt based on intent
# Function to generate reformulated prompts
def reformulate_prompt(user_input, intent_label):
    """
    Reformulate the prompt based on the classified intent.
    """
    core_sentence = extract_core_sentence(user_input)
    prompt_templates = {
        "grammar correction": f"Fix the grammar in this sentence: {core_sentence}",
        "information request": f"Provide information about: {core_sentence}",
        "dialog continuation": f"Continue the conversation based on the previous dialog:\n{core_sentence}\n",
        "personal opinion": f"What is your personal opinion on: {core_sentence}?",
        "product inquiry": f"Provide details about the product: {core_sentence}",
        "feedback request": f"Please provide feedback on: {core_sentence}",
        "recommendation request": f"Recommend something related to: {core_sentence}",
        "clarification request": f"Clarify the following: {core_sentence}",
        "affirmation or agreement": f"Affirm or agree with the statement: {core_sentence}",
    }
    return prompt_templates.get(intent_label, "Input does not require a defined action.")

chat_history = [
            ("Hi there, how are you?", "I am fine. How are you?"),
            ("Tell me a joke!", "The capital of France is Paris."),
            ("Can you tell me another joke?", "Why don't scientists trust atoms? Because they make up everything!"),
            ]
    
    
def generate_response(user_input, model_id):
    try:
        model, tokenizer = get_model_and_tokenizer(model_id)
        device = accelerator.device  # Get the device from the accelerator
       
        # Append chat history
        func_caller = []

        for msg in chat_history:
            func_caller.append({"role": "user", "content": f"{str(msg[0])}"})
            func_caller.append({"role": "assistant", "content": f"{str(msg[1])}"})
 
        highest_label_result = classify_intent(user_input) 

        # Reformulated prompt based on intent classification
        reformulated_prompt = reformulate_prompt(user_input, highest_label_result)
          
        func_caller.append({"role": "user", "content": f'{reformulated_prompt}'})
        formatted_prompt = "\n".join([f"{m['role']}: {m['content']}" for m in func_caller])

        #prompt = user_input
        #device = accelerator.device  # Automatically uses GPU or CPU based on accelerator setup
        
        generation_config = GenerationConfig(
            do_sample=(highest_label == "dialog continuation" or highest_label == "recommendation request"),  # True if dialog continuation, else False
            temperature=0.7 if highest_label == "dialog continuation" else (0.2 if highest_label == "recommendation request" else None),  # Set temperature for specific intents 
            top_k = 5 if highest_label == "recommendation request" else None, 
            #attention_mask=attention_mask,
            max_length=150, 
            repetition_penalty=1.2, 
            length_penalty=1.0, 
            no_repeat_ngram_size=2, 
            num_return_sequences=1,
            pad_token_id=tokenizer.eos_token_id,
            #stop_sequences=["User:", "Assistant:", "\n"],
            )
        
        # Generate response
        gpt_inputs = tokenizer(formatted_prompt, return_tensors="pt").to(device)
        gpt_output = model.generate(gpt_inputs["input_ids"], max_new_tokens=50, generation_config=generation_config)
        final_response = tokenizer.decode(gpt_output[0], skip_special_tokens=True)
        # Extract AI's response only (omit the prompt)
        #ai_response2 = final_response.replace(reformulated_prompt, "").strip()
        ai_response = re.sub(re.escape(formatted_prompt), "", final_response, flags=re.IGNORECASE).strip()
        #ai_response = re.split(r'(?<=\w[.!?]) +', ai_response)
        ai_response = [s.strip() for s in re.split(r'(?<=\w[.!?]) +', ai_response) if s]
       
        # Encode the prompt and candidates
        prompt_embedding = bertmodel.encode(formatted_prompt, convert_to_tensor=True)
        candidate_embeddings = bertmodel.encode(ai_response, convert_to_tensor=True)
        
        # Compute similarity scores between prompt and each candidate
        similarities = util.pytorch_cos_sim(prompt_embedding, candidate_embeddings)[0]
        
        # Find the candidate with the highest similarity score
        
        best_index = similarities.argmax()
        best_response = ai_response[best_index]
        
        # Assuming best_response is already defined and contains the generated response
        
        if highest_label == "dialog continuation":
            # Split the response into sentences
            sentences = best_response.split('. ')
            # Take the first three sentences and join them back together
            best_response = '. '.join(sentences[:3]) if len(sentences) > 3 else best_response
 
        # Append the user's message to the chat history
        chat_history.append({'role': 'user', 'content': user_input})
        chat_history.append({'role': 'assistant', 'content': best_response})

        return best_response
        
    except Exception as e:
        print("Error in generate_response:")
        print(traceback.format_exc())  # Logs the full traceback
        raise e
 
@app.route("/send_message", methods=["POST"])
def handle_post_request():
    try:
        data = request.get_json()
        if data is None:
            return jsonify({"error": "No JSON data provided"}), 400

        message = data.get("inputs", "No message provided.")
        model_id = data.get("model_id", "openai-community/gpt2-large")

        print(f"Processing request with model_id: {model_id}")
        model_response = generate_response(message, model_id)

        return jsonify({
            "received_message": model_response,
            "model_id": model_id,
            "status": "POST request successful!"
        })
    
    except Exception as e:
        print("Error handling POST request:")
        print(traceback.format_exc())  # Logs the full traceback
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)