khurrameycon commited on
Commit
b405fea
·
verified ·
1 Parent(s): 384cf1f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +35 -0
app.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, HTTPException
2
+ from pydantic import BaseModel
3
+ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
4
+
5
+ # Define the input schema
6
+ class ModelInput(BaseModel):
7
+ prompt: str
8
+ max_new_tokens: int = 50 # Optional: Defaults to 50 tokens
9
+
10
+ # Initialize FastAPI app
11
+ app = FastAPI()
12
+
13
+ # Load your model and tokenizer
14
+ model_path = "khurrameycon/SmolLM-135M-Instruct-qa_pairs_converted.json-25epochs" # Update with your model directory
15
+ tokenizer = AutoTokenizer.from_pretrained(model_path)
16
+ model = AutoModelForCausalLM.from_pretrained(model_path)
17
+
18
+ # Initialize the pipeline
19
+ generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
20
+
21
+ @app.post("/generate")
22
+ def generate_text(input: ModelInput):
23
+ try:
24
+ result = generator(
25
+ input.prompt,
26
+ max_new_tokens=input.max_new_tokens,
27
+ return_full_text=False,
28
+ )
29
+ return {"generated_text": result[0]["generated_text"]}
30
+ except Exception as e:
31
+ raise HTTPException(status_code=500, detail=str(e))
32
+
33
+ @app.get("/")
34
+ def root():
35
+ return {"message": "Welcome to the Hugging Face Model API!"}