Update app.py

app.py

@@ -1,85 +1,27 @@
-from fastapi import FastAPI
-import json
-import difflib
-from flair.models import SequenceTagger
-from flair.data import Sentence
+from fastapi import FastAPI, Request
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 app = FastAPI()
 
-# Workaround för PyTorch 2.6 weights_only issue
-original_load = torch.load
-def patched_load(*args, **kwargs):
-    kwargs["weights_only"] = False
-    return original_load(*args, **kwargs)
+# Ladda modellen
+model_id = "AI-Sweden/gpt-sw3-126m"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id)
 
-# Ladda Flair flerspråkig NER-modell
-try:
-    torch.load = patched_load
-    tagger = SequenceTagger.load("flair/ner-multi")
-    torch.load = original_load  # Återställ original torch.load
-except Exception as e:
-    print(f"Error loading model: {str(e)}")
-    raise e
+# Om du kör på CPU – lägg till detta
+device = torch.device("cpu")
+model.to(device)
 
-# Ladda entiteter från entities.json
-with open("entities.json") as f:
-    entities = json.load(f)
-ITEMS = set(entities["items"])
-COLORS = set(entities["colors"])
-PRICES = set(entities["prices"])
+# Input-modell
+class Prompt(BaseModel):
+    text: str
+    max_new_tokens: int = 50
 
-def correct_spelling(word, valid_words, threshold=0.7):
-    """Korrigera stavfel."""
-    normalized = word.rstrip("etn")
-    matches = difflib.get_close_matches(normalized, valid_words, n=1, cutoff=threshold)
-    return matches[0] if matches else word
-
-@app.post("/parse")
-async def parse_user_request(request: str):
-    if not request or len(request) > 200:
-        return {"error": "Ogiltig eller för lång begäran"}
-    try:
-        # Skapa Flair Sentence
-        sentence = Sentence(request)
-        
-        # Prediktera NER-taggar
-        tagger.predict(sentence)
-        
-        # Extrahera entiteter
-        result_entities = {}
-        # Kolla färger och priser i hela meningen
-        words = request.lower().split()
-        for word in words:
-            if word in COLORS:
-                result_entities["FÄRG"] = word
-            elif word in PRICES:
-                result_entities["PRIS"] = word
-        
-        # Extrahera varor från NER
-        for entity in sentence.get_spans("ner"):
-            if entity.tag in ["MISC", "ORG", "LOC"]:  # Diverse, organisationer, platser som potentiella objekt
-                corrected = correct_spelling(entity.text.lower(), ITEMS)
-                if corrected in ITEMS:
-                    result_entities["VARA"] = corrected
-                elif not result_entities.get("VARA"):
-                    result_entities["VARA"] = entity.text.lower()
-        
-        # Om ingen vara hittades
-        if "VARA" not in result_entities:
-            return {"result": "error:ingen vara"}
-        
-        # Skapa strukturerad sträng
-        result_parts = [f"vara:{result_entities['VARA']}"]
-        if "FÄRG" in result_entities:
-            result_parts.append(f"färg:{result_entities['FÄRG']}")
-        if "PRIS" in result_entities:
-            result_parts.append(f"pris:{result_entities['PRIS']}")
-        
-        return {"result": ",".join(result_parts)}
-    except Exception as e:
-        return {"error": f"Fel vid parsning: {str(e)}"}
-
-@app.get("/")
-async def root():
-    return {"message": "Request Parser API is running!"}
+@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}