Update app.py

app.py

@@ -4,6 +4,7 @@ import torch
 
 app = FastAPI()
 
+# Load GPT-SW3-126M
 model_name = "AI-Sweden-Models/gpt-sw3-126m"
 try:
     tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -17,14 +18,26 @@ async def parse_arithmetic_prompt(prompt: str):
     if not prompt or len(prompt) > 100:
         return {"error": "Ogiltig eller för lång prompt"}
     try:
-        messages = [
-            {"role": "system", "content": "Du är en assistent som omvandlar aritmetiska frågor på svenska till strängar, t.ex. '5+5'. Svara endast med strängen."},
-            {"role": "user", "content": prompt}
-        ]
-        input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt")
+        # Format prompt as plain text per GPT-SW3 documentation
+        full_prompt = f"Omvandla följande aritmetiska fråga på svenska till en sträng, t.ex. '5+5'. Fråga: {prompt} Svar: "
+        
+        # Tokenize prompt
+        input_ids = tokenizer(full_prompt, return_tensors="pt").input_ids
+        
+        # Generate response
         with torch.no_grad():
-            output = model.generate(input_ids, max_new_tokens=10, do_sample=False)
+            output = model.generate(
+                input_ids,
+                max_new_tokens=10,
+                do_sample=False,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id
+            )
+        
+        # Decode and extract result
         result = tokenizer.decode(output[0], skip_special_tokens=True).strip()
+        result = result.split("Svar: ")[-1].strip()
+        
         return {"result": result}
     except Exception as e:
         return {"error": f"Fel vid parsning: {str(e)}"}