assignment / app.py
Aroojzahra908's picture
Update app.py
6857c22 verified
Raw
History Blame Contribute Delete
1.24 kB
import os
# ✅ Redirect HuggingFace cache to local folder
os.environ["TRANSFORMERS_CACHE"] = "./cache"
os.environ["HF_HOME"] = "./cache" # just in case huggingface_hub also needs it
from transformers import AutoTokenizer, AutoModelForCausalLM
model_id = "Qwen/Qwen2.5-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
device = 0 if torch.cuda.is_available() else -1
generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
@app.post("/generate")
async def generate_text(request: Request):
data = await request.json()
user_input = data.get("prompt", "").strip()
prompt = f"""
You are an expert academic writer.
Task:
{user_input}
Instructions:
- Write clearly and professionally.
- Avoid repeating the question.
- Use academic tone.
- Format as a structured according to user requirements.
- Do not hallucinate facts.
"""
result = generator(prompt, max_new_tokens=500, do_sample=False)[0]["generated_text"]
if result.lower().startswith(prompt.lower()):
result = result[len(prompt):].strip()
return JSONResponse(content={"result": result})