afnan k salal commited on
Commit
344585f
·
1 Parent(s): 1462422

Add application file

Browse files
Files changed (1) hide show
  1. app.py +18 -1
app.py CHANGED
@@ -1,8 +1,25 @@
1
  from fastapi import FastAPI
 
 
2
 
3
  app = FastAPI()
4
 
 
 
 
 
 
 
 
5
  @app.get("/")
6
  def greet_json():
7
- return {"Hello": "World!"}
 
 
 
 
 
 
 
 
8
 
 
1
  from fastapi import FastAPI
2
+ import torch
3
+ from transformers import AutoModelForCausalLM, AutoTokenizer
4
 
5
  app = FastAPI()
6
 
7
+ # Define model-related constants
8
+ MODEL_NAME = "arnir0/Tiny-LLM"
9
+
10
+ # Global variables to store the tokenizer and model
11
+ tokenizer = None
12
+ model = None
13
+
14
  @app.get("/")
15
  def greet_json():
16
+ global tokenizer, model
17
+
18
+ # Load the model and tokenizer if not already loaded
19
+ if model is None or tokenizer is None:
20
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
21
+ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
22
+ model.eval() # Set model to evaluation mode (optional for inference)
23
+
24
+ return {"Hello": "World!", "model_status": "Loaded and hibernated!"}
25