WhotookNima commited on
Commit
308056c
·
verified ·
1 Parent(s): e37809e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +34 -0
app.py CHANGED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+ import torch
4
+
5
+ # Skapa FastAPI-instans
6
+ app = FastAPI()
7
+
8
+ # Ladda GPT-SW3-126M
9
+ model_name = "AI-Sweden-Models/gpt-sw3-126m"
10
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
11
+ model = AutoModelForCausalLM.from_pretrained(model_name)
12
+
13
+ # Endpoint för att parsa aritmetiska prompts
14
+ @app.post("/parse")
15
+ async def parse_arithmetic_prompt(prompt: str):
16
+ if not prompt or len(prompt) > 100: # Begränsa promptlängd
17
+ return {"error": "Ogiltig eller för lång prompt"}
18
+ try:
19
+ messages = [
20
+ {"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."},
21
+ {"role": "user", "content": prompt}
22
+ ]
23
+ input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt")
24
+ with torch.no_grad():
25
+ output = model.generate(input_ids, max_new_tokens=10, do_sample=False)
26
+ result = tokenizer.decode(output[0], skip_special_tokens=True).strip()
27
+ return {"result": result}
28
+ except Exception as e:
29
+ return {"error": f"Fel vid parsning: {str(e)}"}
30
+
31
+ # Hälsosida för att verifiera att Space körs
32
+ @app.get("/")
33
+ async def root():
34
+ return {"message": "Arithmetic Parser API is running!"}