Update app.py

app.py

@@ -4,6 +4,9 @@ from transformers import AutoTokenizer, AutoModelForQuestionAnswering
 import torch
 import logging
 import sys
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from fastapi.middleware.cors import CORSMiddleware
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -41,9 +44,8 @@ def clean_answer(answer):
     cleaned_answer = ' '.join(token for token in answer.split() if token not in special_tokens)
     return cleaned_answer.strip()
 
-def answer_question(question, system_prompt, temperature, max_new_tokens, top_p, frequency_penalty, presence_penalty, top_k, echo, best_of):
+def answer_question(question):
     logger.info(f"Received question: {question}")
-    logger.info(f"Parameters: temp={temperature}, max_tokens={max_new_tokens}, top_p={top_p}, freq_penalty={frequency_penalty}, pres_penalty={presence_penalty}, top_k={top_k}, echo={echo}, best_of={best_of}")
 
     try:
         logger.info("Combining text from dataset")
@@ -69,36 +71,69 @@ def answer_question(question, system_prompt, temperature, max_new_tokens, top_p,
             logger.warning("Generated answer was empty after cleaning")
             answer = "I'm sorry, but I couldn't find a specific answer to that question based on the Bhagavad Gita. Could you please rephrase your question or ask about one of the core concepts like dharma, karma, bhakti, or the different types of yoga discussed in the Gita?"
 
-        disclaimer = "\n\nPlease note: This response is generated by an AI model based on the Bhagavad Gita. For authoritative information, please consult the original text or scholarly sources."
-        full_response = answer + disclaimer
         logger.info("Answer generated successfully")
 
-        return full_response
+        return answer
 
     except Exception as e:
         logger.error(f"Error in answer_question function: {str(e)}")
         return "I'm sorry, but an error occurred while processing your question. Please try again later."
 
-logger.info("Setting up Gradio interface")
+# FastAPI setup
+app = FastAPI()
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allows all origins
+    allow_credentials=True,
+    allow_methods=["*"],  # Allows all methods
+    allow_headers=["*"],  # Allows all headers
+)
+
+class Question(BaseModel):
+    messages: list
+
+@app.post("/predict")
+async def predict(question: Question):
+    try:
+        last_user_message = next((msg for msg in reversed(question.messages) if msg['role'] == 'user'), None)
+        
+        if not last_user_message:
+            raise HTTPException(status_code=400, detail="No user message found")
+
+        user_question = last_user_message['content']
+
+        answer = answer_question(user_question)
+
+        disclaimer = "\n\nPlease note: This response is generated by an AI model based on the Bhagavad Gita. For authoritative information, please consult the original text or scholarly sources."
+        full_response = answer + disclaimer
+
+        return {"response": full_response, "isTruncated": False}
+
+    except Exception as e:
+        logger.error(f"Error in predict function: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+# Gradio interface (optional, for testing)
 iface = gr.Interface(
     fn=answer_question,
-    inputs=[
-        gr.Textbox(lines=2, placeholder="Enter your question here..."),
-        gr.Textbox(lines=2, placeholder="System prompt (optional)"),
-        gr.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label="Temperature"),
-        gr.Slider(minimum=1, maximum=500, step=1, value=250, label="Max new tokens"),
-        gr.Slider(minimum=0, maximum=1, step=0.05, value=0.95, label="Top p"),
-        gr.Slider(minimum=0, maximum=2, step=0.1, value=1, label="Frequency penalty"),
-        gr.Slider(minimum=0, maximum=2, step=0.1, value=1, label="Presence penalty"),
-        gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Top k"),
-        gr.Checkbox(label="Echo"),
-        gr.Slider(minimum=1, maximum=5, step=1, value=1, label="Best of")
-    ],
+    inputs=gr.Textbox(lines=2, placeholder="Enter your question here..."),
     outputs="text",
     title="Bhagavad Gita Q&A",
     description="Ask a question about the Bhagavad Gita, and get an answer based on the dataset."
 )
 
-logger.info("Launching the Gradio app")
-iface.launch()
-logger.info("Gradio app launched successfully")
+# Run both FastAPI and Gradio
+if __name__ == "__main__":
+    import uvicorn
+    import threading
+    import nest_asyncio
+
+    nest_asyncio.apply()
+
+    def run_fastapi():
+        uvicorn.run(app, host="0.0.0.0", port=7860)
+
+    threading.Thread(target=run_fastapi, daemon=True).start()
+    iface.launch()