Suck it bitches

app.py

@@ -1,31 +1,42 @@
 import gradio as gr
-from transformers import pipeline, AutoModelForQuestionAnswering, AutoTokenizer
+from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
 
-# Load the pre-trained model and tokenizer
-model_name = "distilbert-base-uncased-distilled-squad"
-model = AutoModelForQuestionAnswering.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# 1. Choose a bilingual or multilingual QA model
+MODEL_NAME = "mrm8488/xlm-roberta-large-finetuned-squadv2"
 
-# Load custom knowledge document
-with open("knowledge.txt", "r") as file:
-    custom_knowledge = file.read()
+# 2. Load model + tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForQuestionAnswering.from_pretrained(MODEL_NAME)
 
-# Initialize the question-answering pipeline
+# 3. Initialize QA pipeline
 qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)
 
+# 4. Load or define custom knowledge base
+with open("knowledge.txt", "r", encoding="utf-8") as f:
+    knowledge_text = f.read()
+
+# 5. Define function to answer questions
 def answer_question(question):
-    result = qa_pipeline(question=question, context=custom_knowledge)
-    return result['answer']
+    if not question.strip():
+        return "Please ask a valid question."
+    try:
+        result = qa_pipeline(question=question, context=knowledge_text)
+        return result["answer"]
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-# Set up the Gradio interface
-interface = gr.Interface(
+# 6. Build Gradio interface
+iface = gr.Interface(
     fn=answer_question,
-    inputs="text",
+    inputs=gr.Textbox(lines=2, placeholder="Enter your question here..."),
     outputs="text",
-    title="Custom Knowledge QA",
-    description="Ask questions based on the custom knowledge document."
+    title="Budtender LLM (Bilingual QA)",
+    description=(
+        "A bilingual Q&A model trained on Spanish and English data. "
+        "Ask your cannabis-related questions here!"
+    )
 )
 
+# 7. Launch app
 if __name__ == "__main__":
-    interface.launch()
-
+    iface.launch()