Update app.py

app.py

@@ -3,12 +3,17 @@ from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
+
 app = FastAPI()
 
 # Load your fine-tuned model and tokenizer
-MODEL_NAME = "aubmindlab/aragpt2-medium" 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+MODEL_NAME = "aubmindlab/aragpt2-medium"
+
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+except Exception as e:
+    raise RuntimeError(f"Failed to load model or tokenizer: {str(e)}")
 
 # Define the general prompt template
 general_prompt_template = """
@@ -48,33 +53,37 @@ def generate_text(request: GenerateRequest):
     المادة = request.المادة
     المستوى = request.المستوى
 
-    if not المادة or not المستوى:
-        raise HTTPException(status_code=400, detail="المادة والمستوى مطلوبان.")
+    if not المادة or not المستوى or not isinstance(المادة, str) or not isinstance(المستوى, str):
+        raise HTTPException(status_code=400, detail="المادة والمستوى مطلوبان ويجب أن يكونا نصًا.")
+
+    try:
+        # Format the prompt with user inputs
+        arabic_prompt = general_prompt_template.format(المادة=المادة, المستوى=المستوى)
 
-    # Format the prompt with user inputs
-    arabic_prompt = general_prompt_template.format(المادة=المادة, المستوى=المستوى)
+        # Tokenize the prompt
+        inputs = tokenizer(arabic_prompt, return_tensors="pt", max_length=512, truncation=True)
 
-    # Tokenize the prompt
-    inputs = tokenizer(arabic_prompt, return_tensors="pt", max_length=512, truncation=True)
+        # Generate text
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs.input_ids,
+                max_length=300,
+                num_return_sequences=1,
+                temperature=0.1,
+                top_p=0.9,
+                do_sample=True,
+            )
 
-    # Generate text
-    with torch.no_grad():
-        outputs = model.generate(
-            inputs.input_ids,
-            max_length=300,  # Adjust as needed
-            num_return_sequences=1,
-            temperature=0.1,  # Adjust for creativity
-            top_p=0.9,  # Adjust for diversity
-            do_sample=True,
-        )
+        # Decode the generated text
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Decode the generated text
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Remove the prompt from the generated text
+        generated_text = generated_text.replace(arabic_prompt, "").strip()
 
-    # Remove the prompt from the generated text
-    generated_text = generated_text.replace(arabic_prompt, "").strip()
+        return {"generated_text": generated_text}
 
-    return {"generated_text": generated_text}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error during text generation: {str(e)}")
 
 @app.get("/")
 def read_root():