Update app.py

app.py

@@ -1,143 +1,136 @@
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
-import pandas as pd
-import os
-import requests
-from transformers import GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer, pipeline
-from io import StringIO
-from fastapi.middleware.cors import CORSMiddleware
-from huggingface_hub import HfFolder
-from tqdm import tqdm
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
+import logging
+import re
 
 app = FastAPI()
 
+# Enable CORS if needed
+from fastapi.middleware.cors import CORSMiddleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],  # In production, restrict this to your frontend URL
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-hf_token = os.getenv('HF_API_TOKEN')
-if not hf_token:
-    raise ValueError("Hugging Face API token is not set. Please set the HF_API_TOKEN environment variable.")
-
-
-# Load GPT-2 model and tokenizer
-tokenizer_gpt2 = AutoTokenizer.from_pretrained('gpt2')
-model_gpt2 = GPT2LMHeadModel.from_pretrained('gpt2')
-
-# Create a pipeline for text generation using GPT-2
-text_generator = pipeline("text-generation", model=model_gpt2, tokenizer=tokenizer_gpt2)
-
-
-# Define prompt template for generating the dataset
-prompt_template = """\
-You are an AI designed exclusively for generating synthetic tabular datasets.
-Task: Generate a synthetic dataset based on the provided description and column names.
-Description: {description}
-Columns: {columns}
-Instructions:
-- Output only the tabular data in valid CSV format.
-- Include the header row followed strictly by the data rows.
-- Do not include any additional text, explanations, comments, or code outside of the CSV data.
-- Ensure that the values for each column are contextually appropriate based on the description.
-- Do not alter the column names or add any new columns.
-- Each row must contain data for all columns without any empty values unless specified.
-- Format Example (do not include this line or the following example in your output):
-Column1,Column2,Column3
-Value1,Value2,Value3
-Value4,Value5,Value6
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+####################################
+# Text Generation Endpoint
+####################################
+
+TEXT_MODEL_NAME = "aubmindlab/aragpt2-medium"
+text_tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_NAME)
+text_model = AutoModelForCausalLM.from_pretrained(TEXT_MODEL_NAME)
+
+general_prompt_template = """
+أنت الآن نموذج لغة مخصص لتوليد نصوص عربية تعليمية بناءً على المادة والمستوى التعليمي. سيتم إعطاؤك مادة تعليمية ومستوى تعليمي، وعليك إنشاء نص مناسب بناءً على ذلك. النص يجب أن يكون:
+1. واضحًا وسهل الفهم.
+2. مناسبًا للمستوى التعليمي المحدد.
+3. مرتبطًا بالمادة التعليمية المطلوبة.
+4. قصيرًا ومباشرًا.
+### أمثلة:
+1. المادة: العلوم
+   المستوى: الابتدائي
+   النص: النباتات كائنات حية تحتاج إلى الماء والهواء وضوء الشمس لتنمو. بعض النباتات تنتج أزهارًا وفواكه. النباتات تساعدنا في الحصول على الأكسجين.
+2. المادة: التاريخ
+   المستوى: المتوسط
+   النص: التاريخ هو دراسة الماضي وأحداثه المهمة. من خلال التاريخ، نتعلم عن الحضارات القديمة مثل الحضارة الفرعونية والحضارة الإسلامية. التاريخ يساعدنا على فهم تطور البشرية.
+3. المادة: الجغرافيا
+   المستوى: المتوسط
+   النص: الجغرافيا هي دراسة الأرض وخصائصها. نتعلم عن القارات والمحيطات والجبال. الجغرافيا تساعدنا على فهم العالم من حولنا.
+---
+المادة: {المادة}
+المستوى: {المستوى}
+اكتب نصًا مناسبًا بناءً على المادة والمستوى المحددين. ركّز على جعل النص بسيطًا وواضحًا للمستوى المطلوب.
 """
 
+class GenerateTextRequest(BaseModel):
+    المادة: str
+    المستوى: str
 
-# Define generation parameters
-generation_params = {
-    "top_p": 0.90,
-    "temperature": 0.8,
-    "max_new_tokens": 2048,
-    "return_full_text": False,
-    "use_cache": False
-}
-
-def format_prompt(description, columns):
-    prompt = prompt_template.format(description=description, columns=",".join(columns))
-    return prompt
-
-def generate_synthetic_data(description, columns):
-    formatted_prompt = format_prompt(description, columns)
-    payload = {"inputs": formatted_prompt, "parameters": generation_params}
+@app.post("/generate-text")
+def generate_text(request: GenerateTextRequest):
+    if not request.المادة or not request.المستوى:
+        raise HTTPException(status_code=400, detail="المادة والمستوى مطلوبان.")
 
-    # Call Mixtral model to generate data
-    response = requests.post("https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1", 
-                             headers={"Authorization": f"Bearer {hf_token}"}, json=payload)
-    
-    if response.status_code == 200:
-        return response.json()[0]["generated_text"]
-    else:
-        print(f"Error generating data: {response.status_code}, {response.text}")
-        return None
-
-def process_generated_data(csv_data):
     try:
-        # Ensure the data is cleaned and correctly formatted
-        cleaned_data = csv_data.replace('\r\n', '\n').replace('\r', '\n')
-        data = StringIO(cleaned_data)
-
-        # Read the CSV data with specific parameters to handle irregularities
-        df = pd.read_csv(data)
-
-        return df
-    except pd.errors.ParserError as e:
-        print(f"Failed to parse CSV data: {e}")
-        return None
-
-def generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100):
-    data_frames = []
-
-    for _ in tqdm(range(num_rows // rows_per_generation), desc="Generating Data"):
-        generated_data = generate_synthetic_data(description, columns)
+        prompt = general_prompt_template.format(المادة=request.المادة, المستوى=request.المستوى)
+        inputs = text_tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
         
-        if generated_data:
-           
-            df_synthetic = process_generated_data(generated_data)
-
-            if df_synthetic is not None and not df_synthetic.empty:
-                data_frames.append(df_synthetic)
-            else:
-                print("Skipping invalid generation.")
-
-    if data_frames:
-        return pd.concat(data_frames, ignore_index=True)
-    else:
-        print("No valid data frames to concatenate.")
-        return pd.DataFrame(columns=columns)
-
-class DataGenerationRequest(BaseModel):
-    description: str
-    columns: list[str]
-
-@app.post("/generate/")
-def generate_data(request: DataGenerationRequest):
-    description = request.description.strip()
-    columns = [col.strip() for col in request.columns]
-    csv_data = generate_large_synthetic_data(description, columns, num_rows=1000, rows_per_generation=100)
+        with torch.no_grad():
+            outputs = text_model.generate(
+                inputs.input_ids,
+                max_length=300,
+                num_return_sequences=1,
+                temperature=0.7,
+                top_p=0.95,
+                do_sample=True,
+            )
+        
+        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 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)
+
+def extract_answer(context: str) -> str:
+    """Extract the first sentence (or a key phrase) from the context."""
+    sentences = re.split(r'[.!؟]', context)
+    sentences = [s.strip() for s in sentences if s.strip()]
+    return sentences[0] if sentences else context
+
+def get_question(context: str, answer: str) -> str:
+    """Generate a question based on the context and the candidate answer."""
+    text = f"النص: {context} الإجابة: {answer} </s>"
+    text_encoding = qa_tokenizer.encode_plus(text, return_tensors="pt")
     
-    if csv_data.empty:
-        return JSONResponse(content={"error": "No valid data generated"}, status_code=500)
+    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
+        )
     
-    csv_buffer = StringIO()
-    csv_data.to_csv(csv_buffer, index=False)
-    csv_buffer.seek(0)
+    question = qa_tokenizer.decode(generated_ids[0], skip_special_tokens=True).replace("question:", "").strip()
+    return question
 
-    return StreamingResponse(
-        csv_buffer,
-        media_type="text/csv",
-        headers={"Content-Disposition": "attachment; filename=generated_data.csv"}
-    )
+class GenerateQARequest(BaseModel):
+    text: str
+
+@app.post("/generate-qa")
+def generate_qa(request: GenerateQARequest):
+    if not request.text:
+        raise HTTPException(status_code=400, detail="Text is required.")
+
+    try:
+        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 greet_json():
-    return {"Hello": "World!"}
+    return {"message": "Welcome to the Arabic Text Generation API!"}