Update app.py

app.py

@@ -1,29 +1,29 @@
 from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 import uvicorn
 
 app = FastAPI()
 
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
-# TinyLlama load
+# Load Model
 model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")
 
 @app.post("/api/chat")
 async def chat(request: Request):
     data = await request.json()
     user_input = data.get("prompt", "")
-    system_rules = data.get("system_prompt", "You are Olivia.")
+    system_rules = data.get("system_prompt", "You are a fast AI.")
     
-    # Speed Optimization: Tokens limit 150 (Fast Reply)
-    full_prompt = f"<|system|>\n{system_rules}. System by MT.<|end|>\n<|user|>\n{user_input}<|end|>\n<|assistant|>"
+    prompt = f"<|system|>\n{system_rules}. System by Muhammad Taqi.<|end|>\n<|user|>\n{user_input}<|end|>\n<|assistant|>"
     
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
-    output = model.generate(**inputs, max_new_tokens=150)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    # Speed Optimization: No sampling, fixed token limit
+    output = model.generate(**inputs, max_new_tokens=80, do_sample=False, use_cache=True)
     
     response = tokenizer.decode(output[0], skip_special_tokens=True).split("<|assistant|>")[-1].strip()
     return {"response": response}