Update app.py

app.py

@@ -1,71 +1,69 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+# Load the new model and tokenizer
+model_name = "wedo2910/research_ai"
+tokenizer_name = "wedo2910/research_ai_tok"
 
-model_name = "wedo2910/qa_arabic_model"
-tokenizer = AutoTokenizer.from_pretrained("aubmindlab/bert-base-arabertv02")
-model = AutoModelForQuestionAnswering.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
-qa_pipeline = pipeline(
-    "question-answering",
-    model=model,
-    tokenizer=tokenizer
-)
+# Define the custom inference function
+def single_inference(question, max_new_tokens, temperature):
+    # Prepare the prompt messages
+    messages = [
+        {"role": "system", "content": "اجب علي الاتي بالعربي فقط."},
+        {"role": "user", "content": question},
+    ]
+    
+    # Use the tokenizer's chat template functionality
+    input_ids = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        return_tensors="pt"
+    ).to(model.device)
 
-# Default settings
-default_settings = {
-    "max_new_tokens": 512,
-    "temperature": 0.7,
-    "top_p": 0.9,
-    "min_p": 0,
-    "top_k": 0,
-    "repetition_penalty": 1.0,
-    "presence_penalty": 0,
-    "frequency_penalty": 0,
-    "max_answer_len": 50,
-    "doc_stride": 128,
-}
+    # Define terminator tokens (end-of-sequence markers)
+    terminators = [
+        tokenizer.eos_token_id,
+        tokenizer.convert_tokens_to_ids("<|eot_id|>")
+    ]
 
-# Define a default context (e.g., a general knowledge text or topic)
-default_context = """
-التزم بنص السؤال.
-"""
+    # Generate the output
+    outputs = model.generate(
+        input_ids,
+        max_new_tokens=max_new_tokens,
+        eos_token_id=terminators,
+        do_sample=True,
+        temperature=temperature,
+    )
+    
+    # Decode only the newly generated tokens (i.e. skip the prompt)
+    response = outputs[0][input_ids.shape[-1]:]
+    output = tokenizer.decode(response, skip_special_tokens=True)
+    return output
 
 # Streamlit UI
-st.title("Arabic AI Question Answering")
-st.subheader("Ask a question to get an answer.")
+st.title("Arabic AI Research QA")
+st.subheader("Ask a question to get an answer from the research AI model.")
 
-# Input field for the question only
+# Input field for the question
 question = st.text_input("Question", placeholder="Enter your question here...")
 
-# Settings sliders
+# Settings sliders for generation parameters
 st.subheader("Settings")
-max_new_tokens = st.number_input("Max New Tokens", min_value=1, max_value=1000000, value=512)
-temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=0.7, step=0.1)
-top_p = st.slider("Top P", min_value=0.0, max_value=1.0, value=0.9, step=0.1)
-min_p = st.slider("Min P", min_value=0.0, max_value=1.0, value=0.0, step=0.1)
-top_k = st.number_input("Top K", min_value=0, max_value=1000, value=0)
-repetition_penalty = st.slider("Repetition Penalty", min_value=0.01, max_value=5.0, value=1.0, step=0.1)
-presence_penalty = st.slider("Presence Penalty", min_value=-2.0, max_value=2.0, value=0.0, step=0.1)
-frequency_penalty = st.slider("Frequency Penalty", min_value=-2.0, max_value=2.0, value=0.0, step=0.1)
-max_answer_len = st.number_input("Max Answer Length", min_value=1, value=50)
-doc_stride = st.number_input("Document Stride", min_value=1, value=128)
+max_new_tokens = st.number_input("Max New Tokens", min_value=1, max_value=1000, value=256)
+temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=0.4, step=0.1)
 
 # Generate Answer button
 if st.button("Get Answer"):
     if not question:
         st.error("The question field is required.")
     else:
-        # Generate answer using the default context
         try:
-            prediction = qa_pipeline(
-                {"context": default_context, "question": question},
-                max_answer_len=max_answer_len,
-                doc_stride=doc_stride
-            )
+            answer = single_inference(question, max_new_tokens, temperature)
             st.subheader("Result")
             st.write(f"**Question:** {question}")
-            st.write(f"**Answer:** {prediction['answer']}")
+            st.write(f"**Answer:** {answer}")
         except Exception as e:
-            st.error(f"Error: {e}")
+            st.error(f"Error: {e}")