app.py

'''
from fastapi import FastAPI, Query
from pydantic import BaseModel
import cloudscraper
from bs4 import BeautifulSoup
from transformers import pipeline
import torch
import re
import os

#os.environ["HF_HOME"] = "/home/user/huggingface"
#os.environ["TRANSFORMERS_CACHE"] = "/home/user/huggingface"

app = FastAPI()

class ThreadResponse(BaseModel):
    question: str
    replies: list[str]

def clean_text(text: str) -> str:
    text = text.strip()
    text = re.sub(r"\b\d+\s*likes?,?\s*\d*\s*replies?$", "", text, flags=re.IGNORECASE).strip()
    return text

@app.get("/scrape", response_model=ThreadResponse)
def scrape(url: str = Query(...)):
    scraper = cloudscraper.create_scraper()
    response = scraper.get(url)

    if response.status_code == 200:
        soup = BeautifulSoup(response.content, 'html.parser')
        comment_containers = soup.find_all('div', class_='post__content')

        if comment_containers:
            question = clean_text(comment_containers[0].get_text(strip=True, separator="\n"))
            replies = [clean_text(comment.get_text(strip=True, separator="\n")) for comment in comment_containers[1:]]
            return ThreadResponse(question=question, replies=replies)
    return ThreadResponse(question="", replies=[])

MODEL_NAME = "microsoft/phi-2"

# Load the text-generation pipeline once at startup
text_generator = pipeline(
    "text-generation",
    model=MODEL_NAME,
    trust_remote_code=True,
    device=0 if torch.cuda.is_available() else -1,  # GPU if available, else CPU
)

class PromptRequest(BaseModel):
    prompt: str

@app.post("/generate")
async def generate_text(request: PromptRequest):
    # The model expects a string prompt, so pass request.prompt directly
    outputs = text_generator(
        request.prompt,
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        num_return_sequences=1,
    )
    
    generated_text = outputs[0]['generated_text']

    # Optional: parse reasoning and content if your model uses special tags like </think>
    if "</think>" in generated_text:
        reasoning_content = generated_text.split("</think>")[0].strip()
        content = generated_text.split("</think>")[1].strip()
    else:
        reasoning_content = ""
        content = generated_text.strip()

    return {
        "reasoning_content": reasoning_content,
        "generated_text": content
    }

'''

from fastapi import FastAPI, Query, Path
from pydantic import BaseModel
import cloudscraper
from bs4 import BeautifulSoup
from transformers import AutoTokenizer, AutoModelForCausalLM, T5Tokenizer, T5ForConditionalGeneration
import torch
import re
from fastapi.responses import JSONResponse
from fastapi.requests import Request
from fastapi import status

app = FastAPI()

# --- Data Models ---

class ThreadResponse(BaseModel):
    question: str
    replies: list[str]

class PromptRequest(BaseModel):
    prompt: str

class GenerateResponse(BaseModel):
    reasoning_content: str
    generated_text: str

# --- Utility Functions ---

def clean_text(text: str) -> str:
    text = text.strip()
    text = re.sub(r"\b\d+\s*likes?,?\s*\d*\s*replies?$", "", text, flags=re.IGNORECASE).strip()
    return text

# --- Scraping Endpoint ---

@app.get("/scrape", response_model=ThreadResponse)
def scrape(url: str):
    scraper = cloudscraper.create_scraper()
    response = scraper.get(url)

    if response.status_code == 200:
        soup = BeautifulSoup(response.content, "html.parser")
        comment_containers = soup.find_all("div", class_="post__content")

        if comment_containers:
            question = clean_text(comment_containers[0].get_text(strip=True, separator="\n"))
            replies = [clean_text(comment.get_text(strip=True, separator="\n")) for comment in comment_containers[1:]]
            return ThreadResponse(question=question, replies=replies)
    return ThreadResponse(question="", replies=[])

# --- Load DeepSeek-R1-Distill-Qwen-1.5B Model & Tokenizer ---

deepseek_model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
deepseek_tokenizer = AutoTokenizer.from_pretrained(deepseek_model_name)
deepseek_model = AutoModelForCausalLM.from_pretrained(deepseek_model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
deepseek_model = deepseek_model.to(device)

# --- Load T5-Large Model & Tokenizer ---

t5_model_name = "google-t5/t5-large"
t5_tokenizer = T5Tokenizer.from_pretrained(t5_model_name)
t5_model = T5ForConditionalGeneration.from_pretrained(t5_model_name)
t5_model = t5_model.to(device)

# --- Generation Functions ---

def generate_deepseek(prompt: str) -> (str, str):
    inputs = deepseek_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(device)
    outputs = deepseek_model.generate(
        **inputs,
        max_length=512,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        num_return_sequences=1,
        pad_token_id=deepseek_tokenizer.eos_token_id,
    )
    generated_text = deepseek_tokenizer.decode(outputs[0], skip_special_tokens=True)

    if "</think>" in generated_text:
        reasoning_content, content = generated_text.split("</think>", 1)
        return reasoning_content.strip(), content.strip()
    else:
        return "", generated_text.strip()

def generate_t5(prompt: str) -> (str, str):
    inputs = t5_tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True).to(device)
    outputs = t5_model.generate(
        inputs,
        max_length=512,
        num_beams=4,
        repetition_penalty=2.5,
        length_penalty=1.0,
        early_stopping=True,
    )
    generated_text = t5_tokenizer.decode(outputs[0], skip_special_tokens=True)

    if "</think>" in generated_text:
        reasoning_content, content = generated_text.split("</think>", 1)
        return reasoning_content.strip(), content.strip()
    else:
        return "", generated_text.strip()

# --- API Endpoints ---

@app.post("/generate/{model_name}", response_model=GenerateResponse)
async def generate(
    request: PromptRequest,
    model_name: str = Path(..., description="Model to use: 'deepseekr1-qwen' or 't5-large'")
):
    if model_name == "deepseekr1-qwen":
        reasoning, text = generate_deepseek(request.prompt)
    elif model_name == "t5-large":
        reasoning, text = generate_t5(request.prompt)
    else:
        return GenerateResponse(reasoning_content="", generated_text=f"Error: Unknown model '{model_name}'.")

    return GenerateResponse(reasoning_content=reasoning, generated_text=text)

# --- Global Exception Handler ---

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    print(f"Exception: {exc}")
    return JSONResponse(
        status_code=status.HTTP_200_OK,
        content={
            "reasoning_content": "",
            "generated_text": f"Error: {str(exc)}"
        }
    )