Update app.py

app.py

@@ -1,33 +1,28 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-import os
 
-# Define the path where the model is located
-model_directory = "./tiny-gpt2"  # Make sure this is the correct relative or absolute path to your model folder
-
-# Initialize the model and tokenizer using the correct directory path
-tokenizer = AutoTokenizer.from_pretrained(model_directory)
-model = AutoModelForCausalLM.from_pretrained(model_directory)
-
-# FastAPI app
 app = FastAPI()
 
+# Load model and tokenizer once at startup
+model_name = "FlameF0X/Muffin-2.9b-1C25"  # change this to your own model
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
 class PromptRequest(BaseModel):
     prompt: str
-    max_new_tokens: int = 50  # You can adjust the number of tokens generated
+    max_new_tokens: int = 50
 
 @app.post("/generate")
-async def generate_text(request: PromptRequest):
-    # Encode the input prompt text
-    inputs = tokenizer.encode(request.prompt, return_tensors="pt")
-    
-    # Generate the text using the model
-    with torch.no_grad():
-        outputs = model.generate(inputs, max_length=request.max_new_tokens + len(inputs[0]))
-    
-    # Decode the generated text and return the response
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    return {"generated_text": generated_text}
+async def generate_text(req: PromptRequest):
+    inputs = tokenizer(req.prompt, return_tensors="pt")
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=req.max_new_tokens,
+        do_sample=True,
+        temperature=0.8,
+        top_p=0.95,
+    )
+    generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"generated_text": generated}