Spaces:
Paused
Paused
| import os | |
| import logging | |
| import tempfile | |
| from fastapi import FastAPI, HTTPException | |
| from pydantic import BaseModel | |
| TMP_CACHE = os.environ.get("HF_CACHE_DIR", os.path.join(tempfile.gettempdir(), "hf_cache")) | |
| try: | |
| os.makedirs(TMP_CACHE, exist_ok=True) | |
| except Exception: | |
| TMP_CACHE = tempfile.gettempdir() | |
| os.environ["TRANSFORMERS_CACHE"] = TMP_CACHE | |
| os.environ["HF_HOME"] = TMP_CACHE | |
| os.environ["HF_DATASETS_CACHE"] = TMP_CACHE | |
| os.environ["HF_METRICS_CACHE"] = TMP_CACHE | |
| app = FastAPI(title="DirectEd LoRA API") | |
| class PromptRequest(BaseModel): | |
| prompt: str | |
| max_new_tokens: int = 2048 | |
| def root(): | |
| return {"status": "AI backend is running"} | |
| pipe = None | |
| def load_model(): | |
| global pipe | |
| try: | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline | |
| from peft import PeftModel | |
| BASE_MODEL = "unsloth/llama-3-8b-Instruct-bnb-4bit" | |
| ADAPTER_REPO = "rayymaxx/DirectEd-AI-LoRA" | |
| tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL) | |
| base_model = AutoModelForCausalLM.from_pretrained( | |
| BASE_MODEL, | |
| device_map="auto", | |
| torch_dtype="auto", | |
| ) | |
| model = PeftModel.from_pretrained(base_model, ADAPTER_REPO) | |
| model.eval() | |
| pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto") | |
| logging.info("Model and adapter loaded successfully.") | |
| except Exception as e: | |
| logging.exception("Failed to load model at startup: %s", e) | |
| pipe = None | |
| def generate(req: PromptRequest): | |
| if pipe is None: | |
| raise HTTPException(status_code=503, detail="Model not loaded. Check logs.") | |
| try: | |
| output = pipe(req.prompt, max_new_tokens=req.max_new_tokens, do_sample=True, temperature=0.7) | |
| text = output[0].get("generated_text", "").strip() | |
| if text.startswith(req.prompt): | |
| text = text[len(req.prompt):].strip() | |
| if not text: | |
| text = "No response generated by the model." | |
| return {"response": text} | |
| except Exception as e: | |
| logging.exception("Generation failed for prompt '%s': %s", req.prompt, e) | |
| raise HTTPException(status_code=500, detail=f"Generation failed: {e}") | |