Update app.py

app.py

@@ -1,27 +1,35 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Initialisera modellen och tokenizern
+model_name = "AI-Sweden-Models/gpt-sw3-126m-instruct"
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+model.to(device)
+model.eval()
+
+# FastAPI-applikationen
 app = FastAPI()
 
-# Ladda modellen
-model_id = "AI-Sweden/gpt-sw3-126m"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
+class UserInput(BaseModel):
+    prompt: str
 
-# Om du kör på CPU – lägg till detta
-device = torch.device("cpu")
-model.to(device)
+@app.post("/generate/")
+async def generate_response(user_input: UserInput):
+    prompt = f"<|endoftext|><s>\nUser:\n{user_input.prompt}\n<s>\nBot:"
+    input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(device)
 
-# Input-modell
-class Prompt(BaseModel):
-    text: str
-    max_new_tokens: int = 50
+    generated_token_ids = model.generate(
+        inputs=input_ids,
+        max_new_tokens=100,
+        do_sample=True,
+        temperature=0.6,
+        top_p=1
+    )[0]
 
-@app.post("/generate")
-async def generate_text(prompt: Prompt):
-    inputs = tokenizer(prompt.text, return_tensors="pt").to(device)
-    outputs = model.generate(**inputs, max_new_tokens=prompt.max_new_tokens)
-    generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"response": generated}
+    generated_text = tokenizer.decode(generated_token_ids[len(input_ids[0]):-1], skip_special_tokens=True)
+    return {"response": generated_text.strip()}