Update app.py

app.py

@@ -1,12 +1,9 @@
-# app.py
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+import gradio as gr
 from peft import AutoPeftModelForCausalLM
 from transformers import AutoTokenizer
 import torch
 
-app = FastAPI()
-
+# Load the model and tokenizer
 model_name = "ahmedembedded/AskFAST"
 load_in_4bit = True
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -14,6 +11,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = AutoPeftModelForCausalLM.from_pretrained(model_name, load_in_4bit=load_in_4bit).to(device)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
+# Define the prompt context
 FAST_prompt_context = """You are an admission officer at Fast University Pakistan. Your role is to answer queries related to the admission process at Fast University. You are expected to provide detailed and accurate responses to questions regarding:
 
 - Application deadlines
@@ -44,9 +42,6 @@ Stay focused on admissions-related topics only.
 
 past_prompts = []
 
-class Query(BaseModel):
-    question: str
-
 def get_answer(question: str) -> str:
     if len(past_prompts) >= 10:
         past_prompts.pop(0)
@@ -71,14 +66,8 @@ def get_answer(question: str) -> str:
 
     return response_text.split('Response:')[1].split('### Input:')[0]
 
-@app.post("/predict/")
-async def predict(query: Query):
-    try:
-        answer = get_answer(query.question)
-        return {"answer": answer}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+def predict(question):
+    return get_answer(question)
+
+iface = gr.Interface(fn=predict, inputs="text", outputs="text")
+iface.launch()