app.py

import os
import shutil
import json

# Set custom cache directories to avoid permission issues. New
os.environ["HF_HOME"] = "/tmp/huggingface"
os.makedirs("/tmp/huggingface", exist_ok=True)

os.environ["XDG_CACHE_HOME"] = "/tmp/.cache"
os.makedirs("/tmp/.cache", exist_ok=True)

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
import re

# Hugging Face model config
REPO_NAME = "jaydatech/phi3-finetuned-project"
BASE_MODEL = "microsoft/Phi-3-mini-4k-instruct"
HF_TOKEN = os.getenv("HF_TOKEN")  # Load from Render environment variable

# # Optional: Cleanup if corrupted config is detected
# def check_and_cleanup_corrupt_cache(repo_name: str):
#     cache_dir = os.environ["HF_HOME"]
#     model_dir = os.path.join(cache_dir, f"models--{repo_name.replace('/', '--')}")
#     if os.path.exists(model_dir):
#         for root, dirs, files in os.walk(model_dir):
#             for file in files:
#                 if file == "config.json":
#                     path = os.path.join(root, file)
#                     try:
#                         with open(path, "r") as f:
#                             json.load(f)
#                     except json.JSONDecodeError:
#                         print(f"Corrupted config file detected at {path}, cleaning up...")
#                         shutil.rmtree(model_dir, ignore_errors=True)
#                         return

# check_and_cleanup_corrupt_cache(REPO_NAME)

app = FastAPI()

# Enable CORS for frontend access
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Or specify your frontend domain
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Hugging Face model config
# REPO_NAME = "jaydatech/phi3-finetuned-project"
# BASE_MODEL = "microsoft/Phi-3-mini-4k-instruct"
# HF_TOKEN = os.getenv("HF_TOKEN")  # Load from Render environment variable

config = AutoConfig.from_pretrained(REPO_NAME, token=HF_TOKEN)

device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)

model = AutoModelForCausalLM.from_pretrained(
    REPO_NAME,
    config = config,
    token=HF_TOKEN,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)

# Message and request models
class ChatMessage(BaseModel):
    role: str
    text: str

class ChatRequest(BaseModel):
    message: str
    history: Optional[List[ChatMessage]] = []

def is_farewell(message: str) -> bool:
    farewells = ["bye", "goodbye", "see you", "farewell", "exit", "quit", "end"]
    message_lower = message.lower().strip()
    return any(re.search(rf"\b{re.escape(f)}\b", message_lower) for f in farewells)
    # ============================= Recognizes "attend" but with "end" in it... Fix is above ^ =============================
    # farewells = ["bye", "goodbye", "see you", "farewell", "exit", "quit", "end"]
    # message_lower = message.lower().strip()
    # return any(farewell in message_lower for farewell in farewells)

def clean_input_text(text: str) -> str:
    # Remove any "Instruction N: ..." or similar phrases
    return re.sub(r"Instruction\s*\d+\s*\(.*?\):", "", text, flags=re.IGNORECASE)

@app.post("/chat")
async def chat(request: ChatRequest):
    try:
        history = request.history
        user_message = request.message

        if is_farewell(user_message):
            return {
                "response": "Goodbye! Feel free to chat again if you have more questions.",
                "terminate": True
            }

        conversation = (
            "<|system|>\nYou are an AI assistant for the Federal Reserve Bank of St. Louis. "
            "Answer questions based ONLY on your knowledge of the Federal Reserve Bank of St. Louis. "
            "If the answer is NOT in the training data, respond with: 'I don't think this information is available. Maybe rephrase for me!'. "
            "Answer ONLY what the user asks. Do not volunteer information unless specifically requested. "
            "Do NOT ask: 'How can I assist you today?' or 'What can I do for you today?' after every response you give. "
            "Provide concise answers to the exact question asked and nothing more.\n"
        )


        seen_user_messages = set()
        seen_model_messages = set()
        
        for msg in history:
            if msg.role == "user":
                cleaned_user_text = clean_input_text(msg.text.strip())
                if cleaned_user_text and cleaned_user_text not in seen_user_messages:
                    conversation += f"<|user|>\n{cleaned_user_text}\n"
                    seen_user_messages.add(cleaned_user_text)
        
            elif msg.role == "model":
                cleaned_model_text = clean_input_text(msg.text.strip())
                if cleaned_model_text and cleaned_model_text not in seen_model_messages:
                    conversation += f"<|assistant|>\n{cleaned_model_text}\n"
                    seen_model_messages.add(cleaned_model_text)
        # seen_messages = set()
        # for msg in history:
        #     if msg.role == "user" and msg.text.strip() not in seen_messages:
        #         cleaned_user_text = clean_input_text(msg.text.strip())
        #         conversation += f"<|user|>\n{cleaned_user_text}\n"
        #         seen_messages.add(cleaned_user_text)
        #     elif msg.role == "model":
        #         cleaned_model_text = clean_input_text(msg.text.strip())
        #         conversation += f"<|assistant|>\n{cleaned_model_text}\n"


        #conversation += f"<|user|>\n{user_message.strip()}\n<|assistant|>"
        conversation += f"<|user|>\n{clean_input_text(user_message.strip())}\n<|assistant|>"
        
        inputs = tokenizer(conversation, return_tensors="pt", padding=True, truncation=True, max_length=4096).to(device)

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=130,
                do_sample=True,
                temperature=0.1,
                top_k=5,
                pad_token_id=tokenizer.eos_token_id
            )

        full_response = tokenizer.decode(outputs[0], skip_special_tokens=False)
        assistant_response = ""

        if "<|assistant|>" in full_response:
            assistant_sections = full_response.split("<|assistant|>")
            for section in reversed(assistant_sections):
                cleaned = section.strip()
                if cleaned:
                    cleaned = re.split(
                        r"(<\|user\|>|<\|system\|>|<\|assistant\|>|\nuser[:\s]|<\|endoftext\|>)",
                        cleaned
                    )[0]
                    cleaned = re.sub(r"\n?(User|Assistant)\s*[:：\-–]\s*.*", "", cleaned, flags=re.IGNORECASE).strip()
                    assistant_response = cleaned
                    break

        if not assistant_response:
            assistant_response = "⚠️ Sorry, I couldn't generate a response."

        assistant_response = re.sub(r'\*\* Instruction \*\*:.*?(?=\n\n|\n$|$)', '', assistant_response, flags=re.DOTALL)
        assistant_response = re.sub(r'\*\* Instruction \*\*.*?(?=\n\n|\n$|$)', '', assistant_response, flags=re.DOTALL)
        assistant_response = re.sub(r'\n{3,}', '\n\n', assistant_response).strip()
        assistant_response = re.sub(r"(How can I assist you today\?|What else can I help you with\?|How can I help you today\?)", "", assistant_response, flags=re.IGNORECASE).strip()

        def remove_repeated_sentences(response):
            sentences = response.split(". ")
            seen = set()
            cleaned = []
            for sentence in sentences:
                if sentence not in seen:
                    cleaned.append(sentence)
                    seen.add(sentence)
            return ". ".join(cleaned)

        assistant_response = remove_repeated_sentences(assistant_response)

        return {"response": assistant_response}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))