Update app.py

app.py

@@ -1,3 +1,4 @@
+'''
 from fastapi import FastAPI, Query
 from pydantic import BaseModel
 import cloudscraper
@@ -74,4 +75,111 @@ async def generate_text(request: PromptRequest):
     return {
         "reasoning_content": reasoning_content,
         "generated_text": content
-    }
+    }
+
+'''
+
+from fastapi import FastAPI, Query
+from pydantic import BaseModel
+import cloudscraper
+from bs4 import BeautifulSoup
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+import torch
+import re
+
+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 = 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=[])
+
+
+# --- Load T5-Small Model and Tokenizer ---
+
+tokenizer = T5Tokenizer.from_pretrained("google/t5-small")
+model = T5ForConditionalGeneration.from_pretrained("google/t5-small")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+
+# --- Core Generation Function Using T5 Prompting ---
+
+def generate_text_with_t5(prompt: str) -> (str, str):
+    """
+    Accepts a prompt string that includes the T5 task prefix (e.g. "summarize: ..."),
+    generates output text, and optionally extracts reasoning if present.
+    Returns a tuple (reasoning_content, generated_text).
+    """
+    # Tokenize input prompt with truncation to max 512 tokens
+    inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True).to(device)
+
+    # Generate output tokens with beam search for quality
+    outputs = model.generate(
+        inputs,
+        max_length=512,
+        num_beams=4,
+        repetition_penalty=2.5,
+        length_penalty=1.0,
+        early_stopping=True,
+    )
+
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Optional: parse reasoning if your prompt/model uses a special separator like </think>
+    if "</think>" in generated_text:
+        reasoning_content, content = generated_text.split("</think>", 1)
+        reasoning_content = reasoning_content.strip()
+        content = content.strip()
+    else:
+        reasoning_content = ""
+        content = generated_text.strip()
+
+    return reasoning_content, content
+
+
+# --- /generate Endpoint Using T5 Prompting ---
+
+@app.post("/generate", response_model=GenerateResponse)
+async def generate(request: PromptRequest):
+    """
+    Accepts a prompt string from frontend, which should include the T5 task prefix,
+    e.g. "summarize: {text to summarize}" or "translate English to German: {text}".
+    Returns generated text and optional reasoning content.
+    """
+    reasoning_content, generated_text = generate_text_with_t5(request.prompt)
+    return GenerateResponse(reasoning_content=reasoning_content, generated_text=generated_text)