Batrdj commited on
Commit
454c6b3
·
verified ·
1 Parent(s): b4e15ac

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +30 -0
app.py CHANGED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI
2
+ from pydantic import BaseModel
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM
4
+
5
+ app = FastAPI()
6
+
7
+ # Ultra-tiny model (SAFE for free CPU)
8
+ MODEL_NAME = "sshleifer/tiny-gpt2"
9
+
10
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
11
+ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
12
+
13
+ class Prompt(BaseModel):
14
+ message: str
15
+
16
+ @app.get("/")
17
+ def root():
18
+ return {"status": "TinyLLM API is running"}
19
+
20
+ @app.post("/chat")
21
+ def chat(prompt: Prompt):
22
+ inputs = tokenizer(prompt.message, return_tensors="pt")
23
+ outputs = model.generate(
24
+ **inputs,
25
+ max_new_tokens=50,
26
+ do_sample=True,
27
+ temperature=0.7
28
+ )
29
+ response = tokenizer.decode(outputs[0], skip_special_tokens=True)
30
+ return {"response": response}