Update app.py

app.py

@@ -1,356 +1,24 @@
-'''
-from fastapi import FastAPI, Query
+from fastapi import FastAPI
 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, PegasusTokenizer, PegasusForConditionalGeneration
-import torch
-import re
-from fastapi.responses import JSONResponse
-from fastapi.requests import Request
-from fastapi import status
-from typing import List, Dict, Optional
 from llama_cpp import Llama
 
 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
-
-# New model for summarization request
-class SummarizeRequest(BaseModel):
-    replies: List[str]
-    task: str  # expecting "summarisation"
-
-# New model for summarization response
-class SummarizeResponse(BaseModel):
-    individual_summaries: Dict[int, Dict[str, str]]  # {index: {"reasoning": str, "summary": str}}
-    combined_reasoning: str
-    combined_summary: 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)
-
-pegasus_model_name = "google/pegasus-large"
-pegasus_tokenizer = PegasusTokenizer.from_pretrained(pegasus_model_name)
-pegasus_model = PegasusForConditionalGeneration.from_pretrained(pegasus_model_name)
-pegasus_model = pegasus_model.to(device)
-
-qwen3_model_name = "Qwen/Qwen3-0.6B"
-qwen3_tokenizer = AutoTokenizer.from_pretrained(qwen3_model_name)
-qwen3_model = AutoModelForCausalLM.from_pretrained(qwen3_model_name)
-qwen3_model = qwen3_model.to(device)
-
 qwen3_gguf_llm = Llama.from_pretrained(
     repo_id="unsloth/Qwen3-0.6B-GGUF",
     filename="Qwen3-0.6B-BF16.gguf",
 )
 
+class PromptRequest(BaseModel):
+    prompt: str
 
-# --- 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_new_tokens=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 ---
-
-def generate_pegasus(prompt: str) -> (str, str):
-    # Pegasus expects raw text input (no prefix needed)
-    inputs = pegasus_tokenizer(
-        prompt,
-        return_tensors="pt",
-        truncation=True,
-        max_length=1024,
-    ).to(device)
-
-    outputs = pegasus_model.generate(
-        **inputs,
-        max_new_tokens=150,
-        num_beams=4,
-        length_penalty=2.0,
-        early_stopping=True,
-    )
-    generated_text = pegasus_tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Pegasus does not use <think> tags, so no reasoning extraction
-    return "", generated_text.strip()
-
-def generate_qwen3(prompt: str) -> (str, str):
-    inputs = qwen3_tokenizer(
-        prompt,
-        return_tensors="pt",
-        truncation=True,
-        max_length=1024,
-    ).to(device)
-
-    outputs = qwen3_model.generate(
-        **inputs,
-        max_new_tokens=512,
-        temperature=0.7,
-        top_p=0.9,
-        do_sample=True,
-        num_return_sequences=1,
-        pad_token_id=qwen3_tokenizer.eos_token_id,
-    )
-
-    generated_text = qwen3_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()
+class GenerateResponse(BaseModel):
+    reasoning_content: str = ""
+    generated_text: str
 
-def generate_qwen3_gguf(prompt: str, max_tokens: int = 256) -> (str, str):
-    messages = [
-        {"role": "user", "content": prompt}
-    ]
-    response = qwen3_gguf_llm.create_chat_completion(
-        messages=messages,
-        max_tokens=max_tokens,
-    )
+@app.post("/generate/qwen3-0.6b-gguf", response_model=GenerateResponse)
+async def generate_qwen3_gguf_endpoint(request: PromptRequest):
+    messages = [{"role": "user", "content": request.prompt}]
+    response = qwen3_gguf_llm.create_chat_completion(messages=messages, max_tokens=256)
     generated_text = response['choices'][0]['message']['content']
-    if "</think>" in generated_text:
-        reasoning_content, content = generated_text.split("</think>", 1)
-        return reasoning_content.strip() + "</think>", content.strip()
-    else:
-        return "", generated_text.strip()
-
-# --- New summarization endpoint ---
-
-@app.post("/summarize_thread", response_model=SummarizeResponse)
-async def summarize_thread(request: SummarizeRequest):
-    if request.task.lower() != "summarisation":
-        return JSONResponse(
-            status_code=400,
-            content={"error": "Unsupported task. Only 'summarisation' is supported."}
-        )
-
-    individual_summaries = {}
-    combined_reasonings = []
-    combined_summaries = []
-
-    # Summarize each reply individually
-    for idx, reply in enumerate(request.replies):
-        reasoning, summary = generate_qwen3_gguf(reply, max_tokens=256)
-        individual_summaries[idx] = {
-            "reasoning": reasoning,
-            "summary": summary
-        }
-        if reasoning:
-            combined_reasonings.append(reasoning)
-        combined_summaries.append(summary)
-
-    # Combine all individual summaries into one text
-    combined_summary_text = " ".join(combined_summaries)
-
-    # Recursively summarize combined summary if too long (optional)
-    # Here, we summarize combined summary to get final reasoning and summary
-    final_reasoning, final_summary = generate_qwen3_gguf(combined_summary_text, max_tokens=256)
-
-    # Append final reasoning to combined reasonings
-    if final_reasoning:
-        combined_reasonings.append(final_reasoning)
-
-    return SummarizeResponse(
-        individual_summaries=individual_summaries,
-        combined_reasoning="\n\n".join(combined_reasonings).strip(),
-        combined_summary=final_summary.strip()
-    )
-
-
-
-@app.post("/generate/{model_name}", response_model=GenerateResponse)
-async def generate(
-    request: PromptRequest,
-    model_name: str = Path(..., description="Model to use: 'deepseekr1-qwen', 't5-large', 'pegasus-large', 'qwen3-0.6b-hf', or 'qwen3-0.6b-gguf'")
-):
-    if model_name == "deepseekr1-qwen":
-        reasoning, text = generate_deepseek(request.prompt)
-    elif model_name == "t5-large":
-        reasoning, text = generate_t5(request.prompt)
-    elif model_name == "pegasus-large":
-        reasoning, text = generate_pegasus(request.prompt)
-    elif model_name == "qwen3-0.6b-hf":
-        reasoning, text = generate_qwen3_hf(request.prompt)
-    elif model_name == "qwen3-0.6b-gguf":
-        reasoning, text = generate_qwen3_gguf(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)}"
-        }
-    )
+    return GenerateResponse(generated_text=generated_text)