Update app.py

app.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import torch
 
 app = FastAPI()
@@ -9,18 +9,31 @@ app = FastAPI()
 MODEL_NAME = "nlptown/bert-base-multilingual-uncased-sentiment"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Initialize model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
-classifier = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer, device=DEVICE)
+# Initialize sentiment analysis model
+sentiment_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+sentiment_classifier = pipeline(
+    "sentiment-analysis",
+    model=MODEL_NAME,
+    tokenizer=sentiment_tokenizer,
+    device=DEVICE
+)
+
+# Initialize GPT-2 for text generation
+MODEL_NAME_LARGE = "gpt2-large"
+generation_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_LARGE)
+generation_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME_LARGE).to(DEVICE)
 
 class TextInput(BaseModel):
     text: str
 
+class GenerationInput(BaseModel):
+    prompt: str
+    max_length: int = 100
+
 @app.post("/analyze-sentiment")
 async def analyze_sentiment(input_data: TextInput):
     try:
-        result = classifier(input_data.text)
+        result = sentiment_classifier(input_data.text)
         return {
             "sentiment": result[0]['label'],
             "score": float(result[0]['score'])
@@ -28,31 +41,41 @@ async def analyze_sentiment(input_data: TextInput):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-# Przykład dla większego modelu (np. GPT-2)
-MODEL_NAME_LARGE = "gpt2-large"
-tokenizer_large = AutoTokenizer.from_pretrained(MODEL_NAME_LARGE)
-model_large = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME_LARGE)
-
-class GenerationInput(BaseModel):
-    prompt: str
-    max_length: int = 100
-
 @app.post("/generate-text")
 async def generate_text(input_data: GenerationInput):
     try:
-        inputs = tokenizer_large(input_data.prompt, return_tensors="pt")
-        outputs = model_large.generate(
+        inputs = generation_tokenizer(
+            input_data.prompt, 
+            return_tensors="pt"
+        ).to(DEVICE)
+        
+        outputs = generation_model.generate(
             inputs["input_ids"],
             max_length=input_data.max_length,
             num_return_sequences=1,
-            no_repeat_ngram_size=2
+            no_repeat_ngram_size=2,
+            pad_token_id=generation_tokenizer.eos_token_id
+        )
+        
+        generated_text = generation_tokenizer.decode(
+            outputs[0], 
+            skip_special_tokens=True
         )
-        generated_text = tokenizer_large.decode(outputs[0], skip_special_tokens=True)
+        
         return {"generated_text": generated_text}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-# Dodanie podstawowego health checka
 @app.get("/health")
 async def health_check():
-    return {"status": "healthy"}
+    return {
+        "status": "healthy",
+        "sentiment_model": MODEL_NAME,
+        "generation_model": MODEL_NAME_LARGE,
+        "device": str(DEVICE)
+    }
+
+# Dodaj to na końcu pliku
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)