Upload main.py

backend/app/main.py

@@ -1,22 +1,19 @@
 import os
 from typing import List, Literal, Optional
+import requests
 
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import torch
 
-APP_TITLE = "HF Chat (Fathom-R1-14B)"
+APP_TITLE = "HF Chat (Fathom-R1-14B via API)"
 APP_VERSION = "0.2.0"
 
 MODEL_ID = os.getenv("MODEL_ID", "FractalAIResearch/Fathom-R1-14B")
-PIPELINE_TASK = os.getenv("PIPELINE_TASK", "text-generation")
-MAX_INPUT_TOKENS = int(os.getenv("MAX_INPUT_TOKENS", "8192"))
 STATIC_DIR = os.getenv("STATIC_DIR", "/app/static")
 ALLOWED_ORIGINS = os.getenv("ALLOWED_ORIGINS", "")
-QUANTIZE = os.getenv("QUANTIZE", "auto")  
+HF_API_TOKEN = os.getenv("HF_API_TOKEN")
 
 app = FastAPI(title=APP_TITLE, version=APP_VERSION)
 
@@ -46,121 +43,67 @@ class ChatResponse(BaseModel):
     reply: str
     model: str
 
-tokenizer = None
-model = None
-generator = None
-
-def load_pipeline():
-    global tokenizer, model, generator
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    # Load tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
-    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-    # Determine load strategy
-    load_kwargs = {}
-    dtype = torch.bfloat16 if device == "cuda" else torch.float32
-
-    if device == "cuda":
-        # try quantization if requested
-        if QUANTIZE.lower() in ("4bit", "8bit", "auto"):
-            try:
-                import bitsandbytes as bnb  # noqa: F401
-                if QUANTIZE.lower() == "8bit":
-                    load_kwargs.update(dict(load_in_8bit=True))
-                else:
-                    # 4bit or auto (prefer 4bit)
-                    load_kwargs.update(dict(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16))
-            except Exception:
-                # bitsandbytes not available; fall back to full precision on GPU
-                pass
-        load_kwargs.setdefault("torch_dtype", dtype)
-        load_kwargs.setdefault("device_map", "auto")
-    else:
-        # CPU fallback
-        load_kwargs.setdefault("torch_dtype", dtype)
-
-    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, **load_kwargs)
-
-    generator = pipeline(
-        PIPELINE_TASK,
-        model=model,
-        tokenizer=tokenizer,
-        device_map=load_kwargs.get("device_map", None) or (0 if device == "cuda" else -1),
-    )
-
-@app.on_event("startup")
-def _startup():
-    load_pipeline()
-
 def messages_to_prompt(messages: List[Message]) -> str:
-    """
-    Prefer tokenizer chat template (Qwen-based models ship one). Fallback to a simple transcript.
-    """
-    try:
-        # Convert to HF chat format: list of dicts with role/content
-        chat = [{"role": m.role, "content": m.content} for m in messages]
-        return tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
-    except Exception:
-        # Fallback formatting
-        parts = []
-        for m in messages:
-            if m.role == "system":
-                parts.append(f"System: {m.content}")
-            elif m.role == "user":
-                parts.append(f"User: {m.content}")
-            else:
-                parts.append(f"Assistant: {m.content}")
-        parts.append("Assistant:")
-        return "".join(parts)
-
-def truncate_prompt(prompt: str, max_tokens: int) -> str:
-    ids = tokenizer(prompt, return_tensors="pt", truncation=False)["input_ids"][0]
-    if len(ids) <= max_tokens:
-        return prompt
-    trimmed = ids[-max_tokens:]
-    return tokenizer.decode(trimmed, skip_special_tokens=True)
+    parts = []
+    for m in messages:
+        if m.role == "system":
+            parts.append(f"System: {m.content}")
+        elif m.role == "user":
+            parts.append(f"User: {m.content}")
+        else:
+            parts.append(f"Assistant: {m.content}")
+    parts.append("Assistant:")
+    return "\n".join(parts)
 
 @app.get("/api/health")
 def health():
-    device = next(model.parameters()).device.type if model is not None else "N/A"
-    return {"status": "ok", "model": MODEL_ID, "task": PIPELINE_TASK, "device": device}
+    return {"status": "ok", "model": MODEL_ID, "source": "huggingface-inference-api"}
 
 @app.post("/api/chat", response_model=ChatResponse)
 def chat(req: ChatRequest):
-    if generator is None:
-        raise HTTPException(status_code=503, detail="Model not loaded")
+    if not HF_API_TOKEN:
+        raise HTTPException(status_code=500, detail="HF_API_TOKEN not set")
+
     if not req.messages:
         raise HTTPException(status_code=400, detail="messages cannot be empty")
 
-    raw_prompt = messages_to_prompt(req.messages)
-    prompt = truncate_prompt(raw_prompt, MAX_INPUT_TOKENS)
-
-    gen_kwargs = {
-        "max_new_tokens": req.max_new_tokens,
-        "do_sample": req.temperature > 0,
-        "temperature": req.temperature,
-        "top_p": req.top_p,
-        "repetition_penalty": req.repetition_penalty,
-        "eos_token_id": tokenizer.eos_token_id,
-        "pad_token_id": tokenizer.pad_token_id,
-        "return_full_text": True,
+    prompt = messages_to_prompt(req.messages)
+
+    headers = {
+        "Authorization": f"Bearer {HF_API_TOKEN}"
     }
-    if req.stop:
-        gen_kwargs["stop"] = req.stop
 
-    outputs = generator(prompt, **gen_kwargs)
-    if isinstance(outputs, list) and outputs and "generated_text" in outputs[0]:
-        full = outputs[0]["generated_text"]
+    payload = {
+        "inputs": prompt,
+        "parameters": {
+            "max_new_tokens": req.max_new_tokens,
+            "temperature": req.temperature,
+            "top_p": req.top_p,
+            "repetition_penalty": req.repetition_penalty,
+            "return_full_text": True,
+        }
+    }
+
+    response = requests.post(
+        f"https://api-inference.huggingface.co/models/{MODEL_ID}",
+        headers=headers,
+        json=payload
+    )
+
+    if response.status_code != 200:
+        raise HTTPException(status_code=response.status_code, detail=response.text)
+
+    result = response.json()
+    if isinstance(result, list) and result and "generated_text" in result[0]:
+        full = result[0]["generated_text"]
         reply = full[len(prompt):].strip() if full.startswith(prompt) else full
     else:
-        reply = str(outputs)
+        reply = str(result)
+
     if not reply:
         reply = "(No response generated.)"
+
     return ChatResponse(reply=reply, model=MODEL_ID)
 
-# Serve frontend build (if present)
 if os.path.isdir(STATIC_DIR):
-    app.mount("/", StaticFiles(directory=STATIC_DIR, html=True), name="static")
+    app.mount("/", StaticFiles(directory=STATIC_DIR, html=True), name="static")