Update app.py

app.py

@@ -5,7 +5,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2Se
 import logging
 import re
 
-app = FastAPI(root_path="/")
+app = FastAPI()
 
 # Enable CORS if needed
 from fastapi.middleware.cors import CORSMiddleware
@@ -56,40 +56,37 @@ class GenerateTextRequest(BaseModel):
 
 @app.post("/generate-text")
 def generate_text(request: GenerateTextRequest):
-    المادة = request.المادة
-    المستوى = request.المستوى
-
-    if not المادة or not المستوى:
+    if not request.المادة or not request.المستوى:
         raise HTTPException(status_code=400, detail="المادة والمستوى مطلوبان.")
 
     try:
-        prompt = general_prompt_template.format(المادة=المادة, المستوى=المستوى)
-        inputs = text_tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
+        prompt = general_prompt_template.format(المادة=request.المادة, المستوى=request.المستوى)
+        inputs = text_tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True).to(device)
+        
         with torch.no_grad():
             outputs = text_model.generate(
                 inputs.input_ids,
                 max_length=300,
                 num_return_sequences=1,
-                temperature=0.1,
-                top_p=0.9,
+                temperature=0.7,
+                top_p=0.95,
                 do_sample=True,
             )
-        generated_text = text_tokenizer.decode(outputs[0], skip_special_tokens=True)
-        # Remove the prompt from the generated text
-        generated_text = generated_text.replace(prompt, "").strip()
+        
+        generated_text = text_tokenizer.decode(outputs[0], skip_special_tokens=True).replace(prompt, "").strip()
         logger.info(f"Generated text: {generated_text}")
         return {"generated_text": generated_text}
+    
     except Exception as e:
         logger.error(f"Error during text generation: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Error during text generation: {str(e)}")
 
 ####################################
-# Question & Answer Generation Endpoint
+# Question & Answer Generation Model
 ####################################
-
 QA_MODEL_NAME = "Mihakram/AraT5-base-question-generation"
 qa_tokenizer = AutoTokenizer.from_pretrained(QA_MODEL_NAME)
-qa_model = AutoModelForSeq2SeqLM.from_pretrained(QA_MODEL_NAME)
+qa_model = AutoModelForSeq2SeqLM.from_pretrained(QA_MODEL_NAME).to(device)
 
 def extract_answer(context: str) -> str:
     """Extract the first sentence (or a key phrase) from the context."""
@@ -99,45 +96,48 @@ def extract_answer(context: str) -> str:
 
 def get_question(context: str, answer: str) -> str:
     """Generate a question based on the context and the candidate answer."""
-    text = "النص: " + context + " " + "الإجابة: " + answer + " </s>"
-    text_encoding = qa_tokenizer.encode_plus(text, return_tensors="pt")
-    qa_model.eval()
-    generated_ids = qa_model.generate(
-        input_ids=text_encoding['input_ids'],
-        attention_mask=text_encoding['attention_mask'],
-        max_length=64,
-        num_beams=5,
-        num_return_sequences=1
-    )
-    question = qa_tokenizer.decode(generated_ids[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
-    # Optionally remove a leading phrase if present
-    question = question.replace('question: ', '').strip()
+    text = f"النص: {context} الإجابة: {answer} </s>"
+    text_encoding = qa_tokenizer.encode_plus(text, return_tensors="pt").to(device)
+    
+    with torch.no_grad():
+        generated_ids = qa_model.generate(
+            input_ids=text_encoding['input_ids'],
+            attention_mask=text_encoding['attention_mask'],
+            max_length=64,
+            num_beams=5,
+            num_return_sequences=1
+        )
+    
+    question = qa_tokenizer.decode(generated_ids[0], skip_special_tokens=True).replace("question:", "").strip()
     return question
 
-def generate_question_answer(context: str):
-    answer = extract_answer(context)
-    question = get_question(context, answer)
-    return question, answer
-
 class GenerateQARequest(BaseModel):
     text: str
 
 @app.post("/generate-qa")
 def generate_qa(request: GenerateQARequest):
-    context = request.text
-    if not context:
+    if not request.text:
         raise HTTPException(status_code=400, detail="Text is required.")
+
     try:
-        question, answer = generate_question_answer(context)
+        question, answer = get_question(request.text, extract_answer(request.text))
         logger.info(f"Generated QA -> Question: {question}, Answer: {answer}")
         return {"question": question, "answer": answer}
+    
     except Exception as e:
         logger.error(f"Error during QA generation: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Error during QA generation: {str(e)}")
 
+####################################
+# Root Endpoint
+####################################
 @app.get("/")
 def read_root():
     return {"message": "Welcome to the Arabic Text Generation API!"}
+
+####################################
+# Running the FastAPI Server
+####################################
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)